PhaseFlow4D: Physically Constrained 4D Beam Reconstruction via Feedback-Guided Latent Diffusion

Abstract: We address the problem of recovering a time-varying 4D distribution from a sparse sequence of 2D projections - analogous to novel-view synthesis from sparse cameras, but applied to the 4D transverse phase space density $ρ(x,p_x,y,p_y)$ of charged particle beams. Direct single shot measurement of this high-dimensional distribution is physically impossible in real particle accelerator systems; only limited 1D or 2D projections are accessible. We propose PhaseFlow4D, a feedback-guided latent diffusion model that reconstructs and tracks the full 4D phase space from incomplete 2D observations alone, with built-in hard physics constraints. Our core technical contribution is a 4D VAE whose decoder generates the full 4D phase space tensor, from which 2D projections are analytically computed and compared against 2D beam measurements. This projection-consistency constraint guarantees physical correctness by construction - not as a soft penalty, but as an architectural prior. An adaptive feedback loop then continuously tunes the conditioning vector of the latent diffusion model to track time-varying distributions online without retraining. We validate on multi-particle simulations of heavy-ion beams at the Facility for Rare Isotope Beams (FRIB), where full physics simulations require $\sim$6 hours on a 100-core HPC system. PhaseFlow4D achieves accurate 4D reconstructions 11000$\times$ faster while faithfully tracking distribution shifts under time-varying source conditions - demonstrating that principled generative reconstruction under incomplete observations transfers robustly beyond visual domains.

• The paper introduces a 4D VAE with hard projection-consistency priors and latent diffusion to reconstruct high-resolution beam phase space from sparse 2D projections.
• It employs conditional latent diffusion with a 20-dimensional operational vector and 100 denoising steps, achieving reconstruction speeds 11000× faster than traditional simulations.
• Adaptive tracking via model-free extremum seeking enables real-time adjustment and accurate recovery of time-varying beam dynamics across all phase space dimensions.

Physically Constrained 4D Beam Phase Space Reconstruction with Feedback-Guided Latent Diffusion

Introduction

"PhaseFlow4D: Physically Constrained 4D Beam Reconstruction via Feedback-Guided Latent Diffusion" (2604.03885) addresses the inverse problem of reconstructing a time-varying 4D transverse phase space density $\rho(x, x', y, y')$ of charged particle beams using only sparse 2D projections. Unlike standard visual scene reconstruction where multiple calibrated images are available, in high-intensity accelerator systems direct measurement of the high-dimensional phase space is infeasible; only a subset of 2D projections is observable. This paper proposes a framework combining VAE-based hard physics priors with feedback-guided conditional latent diffusion, enabling real-time, physically consistent recovery and adaptive tracking of the 4D phase space from minimal 2D data.

Methodological Framework

4D VAE with Architectural Projection Consistency

The core of the PhaseFlow4D system is a 4D VAE: its encoder maps $128^4$ phase space tensors to a $16\times 16\times 4$ latent space, and the decoder reconstructs high-resolution 4D densities. A distinguishing feature is the enforcement of an architectural projection-consistency constraint. All six possible 2D marginal projections are computed analytically from the VAE output, and the training objective includes an $\ell_2$ penalty ensuring that generated projections exactly match the true ones. This guarantees, by construction rather than as a soft regularizer, that all decoded 2D projections originate from a single consistent 4D distribution.

Figure 1: 4D VAE architecture maps high-dimensional phase space tensors to a compact latent code and reconstructs with analytic marginal projections for consistency enforcement.

Conditional Latent Diffusion in Compressed Space

A conditional latent diffusion model (LDM) operates on the VAE latent space, incorporating domain-specific conditioning signals: a 20-dimensional vector encoding beam species, accelerator location, charge neutralization, and relevant magnetic optics settings. This maps operational parameters to valid 4D phase space distributions. The conditional vector is injected through cross-attention layers in a U-Net backbone, and 100 denoising steps are employed for synthesis of the desired latent code.

Figure 2: The conditional latent diffusion employs a U-Net architecture with residual blocks, attention, and 100 denoising steps for latent generation.

During inference, the pipeline is strictly forward: conditional vector $\mathbf{c}$ is mapped via the LDM to a latent embedding, which is then decoded by the VAE, yielding the full 4D density and all required analytic projections.

Figure 3: Visualization of denoising steps in the conditional latent diffusion process for representative beamline conditions.

Adaptive Tracking via Model-Free Extremum Seeking

To enable adaptive tracking of time-varying distributions in practical use, the conditioning vector $\mathbf{c}$ is optimized online using model-free extremum seeking (ES). At each time point, only a single 2D projection measurement $\pi_{xy}(t)$ is available. ES iteratively adjusts $\mathbf{c}$ to minimize the discrepancy between the observed projection and the projection computed from the generated 4D tensor. This gradient-free technique leverages oscillatory dithers to infer descent directions in the cost landscape, circumventing the need for analytical derivatives or Monte Carlo gradients.

Figure 4: Overview of the feedback-driven adaptive framework for real-time tracking of time-varying beam conditions by tuning the conditional vector.

As a consequence of strict projection-consistency, minimization over the observable projection implicitly enforces consistency in all other unmeasured projections, as demonstrated empirically.

Experimental Results

Dataset and Simulation Details

The experimental validation utilizes synthetic data from TRACK-based multi-particle simulations on the FRIB charge selection section, covering 13 isotopes and incorporating realistic machine noise (solenoid and dipole variations, time-varying charge neutralization). Over 300,000 unique $128^4$ phase space tensors and $\sim$1.8 million 2D projections are available for training, validation, and test.

Reconstruction Fidelity and Speed

Quantitative results show high reconstruction fidelity for both training and held-out test distributions. 1D marginal fits ($128^4$0, $128^4$1, $128^4$2, $128^4$3) extracted from generated projections closely match ground truth statistics.

Figure 5: Discrepancy statistics for train/test datasets and detailed projection comparisons for a representative test example.

A single forward pass through the inference pipeline reconstructs the entire $128^4$4 density in about 2 seconds on an H100 GPU, over four orders of magnitude ($128^4$5) faster than the original TRACK-based physics simulation (6 hours on 100 cores).

Robustness to Distribution Drift

Adaptive tracking experiments emulate non-stationary scenarios (e.g., charge neutralization drift). Only the $128^4$6 projection is observed in real time; the ES loop tunes $128^4$7 to minimize projection mismatch. Metrics show that accurate tracking of the observable projection induces accurate recovery of the full 4D phase space, with tight matching in all 1D marginal statistics.

Figure 6: (Top) Tracking of the observable $128^4$8 projection; (Bottom) accurate prediction of all 1D projection statistics throughout the drift process.

Charge state inference via the same mechanism demonstrates that the proposed system recovers beam conditions under substantial unobserved state evolution.

Figure 7: Adaptive tracking of time-varying charge state, showing projection agreement and error with/without tracking.

Detailed comparisons of predicted versus ground truth projections across test examples further corroborate the consistency and fidelity of the approach.

Implications and Future Prospects

This work establishes a paradigm for high-dimensional distribution tracking where only sparse projections are available, shifting from vision-centric generative priors to physics-anchored structural constraints. The strict enforcement of projection consistency differentiates this method from previous hybrid approaches that post-process individual projections, avoiding the disconnect between refinement and physical plausibility.

Practical implications include real-time accelerator beam diagnostics, which historically required expensive and slow high-fidelity simulations. More broadly, the methodology is extensible to any time-varying physical system only accessible through low-dimensional projections, including, but not limited to, medical tomography, plasma diagnostics, and turbulent flow estimation.

Theoretically, PhaseFlow4D demonstrates that LDMs, when correctly structured with hard architectural physics priors, can perform as metric-accurate surrogates for computationally demanding physical simulators, and that inference-time conditional optimization, driven by model-free feedback, can adapt these generative models to dynamic operating landscapes without retraining. This points to a general recipe for embedding invariants and physical symmetries as functional constraints within generative models.

Conclusion

PhaseFlow4D delivers a real-time framework for physically constrained 4D phase space reconstruction and adaptive tracking of charged particle beams, achieving manifold acceleration over direct simulation while guaranteeing strict projection consistency by architectural design. This work extends the scope of generative reconstruction methods beyond the visual domain, signaling future advances in AI-accelerated diagnostics for complex, partially observed physical systems.