BAM! Bayesian Anything Model: a foundation model for generative computational imaging

• Plug-and-play methods use pre-trained priors in iterative likelihood-enforcing schemes, while end-to-end methods typically train physics-aware networks per operator and return point estimates.Zero-shot plug-and-play methods and modern diffusion priors can add biased likelihood approximations, tunable guidance weights, and hundreds of neural function evaluations.Source: §1 · p. 2
• BAM transports a Gaussian reference to the posterior with a conditional flow map, specifying instrument physics at inference rather than fixing it during training.It draws posterior samples in a few steps without likelihood approximation or guidance weights.Source: §1 · p. 2
• BAM’s flow-map training recipe upgrades an unfolded reconstruction network through Lagrangian self-distillation on a normalized stochastic interpolant of the measurement.Source: §1 · p. 2
• Built on the 36M-parameter RAM backbone, BAM learns across operators, noise levels, and datasets and generalizes to new data and tasks.Source: §1 · p. 2
• In three steps, BAM outperforms specialized models and leading zero-shot methods in sample quality across FFHQ, AFHQ, LSUN, DIV2K, and Köhler.The comparison advances the quality–cost Pareto frontier.Source: §1 · p. 2, Figure 2 · p. 2
• BAM is presented as a way to lower barriers to generative computational imaging and reduce the economic and environmental cost of training imaging models.The authors also describe it as opening a new direction for physics-aware Bayesian imaging.Source: §1 · p. 2

• Diffusion models learn to reverse gradual Gaussian corruption, and standard solvers may require hundreds of network evaluations per sample.The noising process forms rescaled, Gaussian-smoothed marginals of the data distribution.Source: §2 · p. 3
• Flow matching learns the probability-flow ODE’s velocity field directly, using least-squares regression on data and noise pairs rather than simulating ODE trajectories.It learns the same ODE as diffusion models, with differences in parametrization, loss weighting, and solver choice.Source: §2 · p. 3
• Consistency models learn the integrated flow that maps a noisy state to its trajectory endpoint, enforcing self-consistency and a boundary condition.They can be distilled from a pre-trained diffusion model or trained directly from data.Source: §2 · p. 3
• Flow maps generalize consistency models by mapping between any two times along a trajectory, with the equal-time map fixed as the identity.The map from time 1 to 0 transports pure noise to data.Source: §2 · p. 3
• The Lagrangian self-distillation loss propagates the flow drift to finite jumps, alongside a flow-matching loss on the diagonal.The times are sampled randomly, and an exponential moving average of the network serves as a detached teacher.Source: §2 · p. 3
• BAM applies this formulation to few-step, physics-aware posterior sampling that transfers to unseen data and tasks zero-shot or with minimal finetuning.Source: §2 · p. 3

• BAM targets posterior samples for imaging problems with known linear operators and noise levels, amortizing one model across randomized operator and noise families.When the forward model changes, the model requires no retraining or at most light finetuning.Source: §3 · p. 4
• BAM conditions its posterior flow map on a rescaled measurement and operator, using a schedule whose noise-to-scaling ratio matches the measurement noise.The flow transports a Gaussian reference distribution to the posterior.Source: §3 · p. 4
• The conditional flow map is parameterized to be the identity when its two time inputs match, while its diagonal velocity follows the conditional probability-flow drift.Off-diagonal backward jumps follow the conditional Lagrangian condition; the noise level enters through measurement rescaling rather than a separate network input.Source: §3 · p. 4
• BAM uses a RAM backbone that embeds unrolled physics-aware updates in a DRUNet, conditioning on augmented measurements and the stacked forward operator.Its two noise-level embeddings are assigned to the flow times.Source: §3 · p. 4
• Training combines a diagonal flow-matching objective with an off-diagonal objective that matches the endpoint derivative to the teacher’s diagonal velocity.The network is trained on interpolated image-noise states with measurements generated using sampled operators.Source: §3 · p. 4
• Sampling pairs each iterate with a measurement interpolant at a matching noise level, then maps it back to time zero at each step.For small measurement noise, the final step uses that noise level; for large noise, the method uses an intermediate level and a Brownian-bridge draw to avoid a long final jump.Source: §3 · p. 5

4.1 SETUP

• Evaluation uses held-out LSUN Bedroom, FFHQ, AFHQ cat and DIV2K images at σy ∈ {0.025, 0.05}; training combines image datasets with six inverse-problem families.The baseline is trained for 170 hours on eight H100 GPUs; BAM* receives 16 additional hours on two H100 GPUs per target dataset while still training across all problems.Source: §4.1 · pp. 5–6
• The evaluation reports PSNR, LPIPS, CMMD and FID, treats CMMD as more reliable than FID on small test sets, and compares zero-shot and training-based baselines.UD2M and I2SB are evaluated only on FFHQ-512; several training-based baselines use one checkpoint per dataset–problem pair, unlike BAM’s single model across noise levels.Source: §4.1 · pp. 5–6

4.2 RESULTS

• Across four datasets, five problems and two noise levels, a BAM variant achieves the best LPIPS in 38 of 40 settings and best CMMD in 36 of 40.The reported exceptions are FFHQ deblurring for LPIPS and one super-resolution setting each on AFHQ and LSUN for CMMD.Source: §4.2 · p. 6
• INT8 TensorRT cuts latency on a 512 × 512 image from 281 ms to 80 ms (3.5×), with PSNR −1.3 dB and LPIPS 0.400 → 0.427.Separately, BAM* advances the baselines’ quality–cost Pareto frontier, with the same conclusion across metrics, inverse problems and datasets.Source: §4.2 · p. 6
• After CT fine-tuning, three-step BAM samples recover fine lung vessels; standard deviation across 64 draws concentrates on anatomical edges and vessels and matches the empirical posterior-mean error’s order of magnitude.This sparse-view CT test uses 51 parallel-beam projections on LIDC-IDRI chest slices, a modality far from the natural images used in pre-training.Source: §4.2 · p. 6
• On the Köhler dataset, BAM outperforms RAM and Carbajal et al.’s original unrolled PnP architecture despite never being trained on spatially varying operators.The kernels are estimated using the pre-trained Kernel Predictor Network; joint fine-tuning with it is proposed as a next step.Source: §4.2 · pp. 6–7
• After fine-tuning on noisy, JPEG-compressed images, BAM generates missing details and regularizes noisy areas; extending it to more general degradation models remains future work.This is a preliminary exploration of nonlinear imaging problems.Source: §4.2 · p. 7
• The empirical mean of repeated BAM draws exceeds RAM*’s PSNR on four of five problems, with super-resolution the exception.The paper characterizes the individual draws as perceptually sharp while their average remains an accurate posterior-mean estimate.Source: §4.2 · p. 8

5 CONCLUSION AND LIMITATIONS

• Across linear inverse problems on FFHQ, AFHQ, LSUN, DIV2K, and Köhler, BAM outperformed specialized and leading zero-shot methods in three steps at a fraction of their computational cost.BAM is a 36M-parameter model that generalizes to unseen data and tasks, zero-shot or with minimal finetuning.Source: §5 · p. 8
• BAM supports additive Gaussian noise and linear or mildly nonlinear operators, requires the instrument model as input, and may return unreliable posteriors without warning under strong distribution shift.Posterior sample quality is assessed empirically; formal guarantees for the learned conditional flow map remain open.Source: §5 · p. 8
• Scientific and medical imaging applications named by the authors include MRI, CT, microscopy, and astronomy, where operators are known and ground truth is scarce.The authors identify posterior samples as useful for Bayesian decision theory, optimal estimation, uncertainty quantification, and experimental design.Source: §5 · p. 8
• The authors identify formal guarantees, non-Gaussian noise, nonlinear or partially unknown operators, larger operator libraries, and out-of-distribution detection as priorities for future work.They also propose more sophisticated backbones and finetuning rather than training from scratch.Source: §5 · p. 9

A.1 METHODS USED IN THE NUMERICAL COMPARISONS

• RAM uses a physics-aware DRUNet with proximal initialization and Krylov features, but its squared-error objective makes it a point estimator targeting E[x | y, A].BAM retains this lightweight backbone but learns a conditional flow map with a random source rather than a deterministic observation-to-image regression.Source: §A.1 · p. 14
• TReg uses text-conditioned latent diffusion with alternating data-consistency optimization and latent updates, requiring hundreds of neural function evaluations.Its physical operator enters inference-time optimization rather than the diffusion backbone.Source: §A.1 · p. 14
• LATINO combines a pretrained latent consistency prior with proximal data-fidelity updates, avoiding differentiation through the generator but still alternating generation and external consistency steps.LATINO requires few neural evaluations; LATINO-PRO additionally estimates text conditioning from the observation through iterative prompt updates, adding calibration computation.Source: §A.1 · p. 14
• UD2M unfolds and distills a diffusion sampling algorithm into a few-step conditional sampler with explicit likelihood-dependent proximal updates.It supports joint training across likelihood families and specialization to the measurement model at inference time.Source: §A.1 · p. 14
• SILO learns a latent degradation surrogate for a pretrained diffusion prior, but its sampling still differentiates through the latent denoiser and requires hundreds of neural evaluations.Unlike BAM’s use of forward and adjoint operator actions, SILO conditions measurements through a learned latent surrogate.Source: §A.1 · pp. 14–15
• I2SB learns a diffusion bridge from paired clean and degraded images, but its standard formulation does not supply an arbitrary forward operator to the network.The comparison adapts an image-to-image bridge to each evaluated restoration task, unlike BAM’s operator-conditioned flow map.Source: §A.1 · p. 15

A.2 RELATED OPERATOR-CONDITIONED AND FLOW-MAP APPROACHES

• InvFusion is an operator-conditioned diffusion sampler focused on strongly underdetermined problems, with a heavier HDiT backbone and a reported 63-NFE configuration.Its benchmarks emphasize settings such as patch inpainting retaining less than 10% of patches, whereas this work focuses on compact sampling across restoration problems where measurement fidelity remains central.Source: §A.2 · p. 15
• FlowMapSR transports paired low-resolution image latents to high-resolution reconstructions, but does not condition on a supplied forward operator or produce posterior distributions for a fixed observation.BAM instead retains a random source while conditioning on the observation and acquisition physics.Source: §A.2 · p. 15

B BAM! ARCHITECTURE AND MODIFICATIONS RELATIVE TO RAM

• RAM initializes reconstruction from A⊤y, refines it with a proximal measurement-consistency step, and supplies forward-operator information through Krylov features.BAM builds on this physics-aware encoder–decoder structure.Source: B BAM! ARCHITECTURE AND MODIFICATIONS… · p. 15
• BAM additionally conditions on the current flow state xt and flow times (s, t), alongside physics-aware features from the measurement yσ.The additional inputs adapt RAM’s conditioning mechanism to the flow formulation.Source: B BAM! ARCHITECTURE AND MODIFICATIONS… · p. 15
• A learned fusion combines measurement and flow-state features using scalar gains and a 1×1 convolution, and the network output is negated to obtain the flow prediction vθ.Apart from this conditioning mechanism, the RAM backbone remains largely unchanged.Source: B BAM! ARCHITECTURE AND MODIFICATIONS… · p. 16

C TRAINING OBJECTIVES AND IMPLEMENTATION DETAILS

• Training samples an operator from a configurable bank per minibatch, with independent Gaussian image and measurement noises and a shared source time drawn uniformly from (10^-4, 1).Restricting the operator bank supports training or finetuning for a particular problem; the conditioning pair uses the noiseless forward projection and equation (7).Source: C TRAINING OBJECTIVES AND… · p. 16
• After 5000 updates, finetuning selects the off-diagonal branch with probability pD = 0.25 when t > δ = 10^-4; otherwise it uses the diagonal branch.The expected post-warm-up objective weights the diagonal main loss by 1 − pD and the off-diagonal LSD loss by pD.Source: C TRAINING OBJECTIVES AND… · p. 16
• LPIPS is introduced after 100 updates with limiting weight 10^-1, while the baseline contrast term starts after 5000 updates and is removed during finetuning.The perceptual loss is evaluated on the clipped clean prediction, and the contrast penalty acts on excessive contrast relative to a reference before clipping.Source: C TRAINING OBJECTIVES AND… · pp. 16–17
• For the LSD objective, a Jacobian–vector product computes the flow map and endpoint derivative jointly, with the teacher output detached and gradients retained through the student derivative.The implementation treats scalar noise levels as fixed when computing the JVP.Source: C TRAINING OBJECTIVES AND… · p. 16

D ABLATION STUDIES

• The ablations modify one architectural or training component at a time while keeping the remaining setup unchanged.The active flow-matching/LSD path uses AdamW with effective learning rate 10^-5, weight decay 10^-3, and gradient-norm clipping at 1.5.Source: D ABLATION STUDIES · p. 17, C TRAINING OBJECTIVES AND… · p. 17

D.1 EFFECT OF NOISE-AWARE CONDITIONING (yσ VS. y)

• 26.23 ± 1.20 dB mean PSNR with yσ versus 22.22 ± 0.60 dB with y on 16 LSDIR images for Gaussian deblurring at σy = 0.05.The authors report an approximately 4.0 dB improvement and visually closer reconstructions with noise-aware conditioning.Source: §D.1 · p. 17, Figure 7 · p. 17, Figure 8 · p. 17

D.2 WHY THE MODIFIED RAM ARCHITECTURE MATTERS

• The modified backbone reaches substantially lower training LPIPS than stacked-operator RAM throughout the first 60 k steps and better preserves fine structure in DIV2K inpainting.It replaces the final stacked conditioning block with separate measurement and image-state branches before fusion.Source: §D.2 · p. 18

D.3 WHY THE CONTRAST LOSS MATTERS

• Without the contrast term, mixed-dataset pre-training recurrently produces over-saturated outputs; the term is removed during dataset-specific fine-tuning.The authors interpret the issue as a possible consequence of heterogeneous image statistics, intensity scales, and operator conditioning ranges.Source: §D.3 · p. 18

D.4 WHY AN ADVERSARIAL LOSS IS NOT REQUIRED

• Adding a discriminator loss during FFHQ super-resolution fine-tuning yields nearly indistinguishable training LPIPS, so the final objective omits it.Omission avoids the additional discriminator, optimization instability, and computational overhead.Source: §D.4 · p. 19

D.5 NUMBER OF SAMPLING STEPS

• Three steps give the highest PSNR in all ten tested settings, although two steps yield lower LPIPS than three at σy = 0.05 for four problems.The four LPIPS exceptions are super-resolution, inpainting, demosaicing, and compressed sensing; each step requires one network evaluation.Source: §D.5 · p. 19, Table 2 · p. 20, Table 3 · p. 20

E.1 QUANTITATIVE COMPARISONS

• Evaluation coverage is constrained: SILO is unavailable for compressed sensing, while UD2M, I2SB, and CoSIGN are evaluated on limited task settings because they require heavy task-specific fine-tuning.The comparison tables cover AFHQ, DIV2K, FFHQ, and LSUN; BAM and BAM⋆ use three steps.Source: §E.1 · p. 23, Table 4 · p. 23, Table 5 · p. 24

E.2 QUALITATIVE COMPARISONS

• The qualitative comparison figures present restorations across AFHQ, DIV2K, FFHQ, and LSUN, with matched regions magnified for inspection.BAM and BAM⋆ use three steps in the AFHQ, DIV2K, and FFHQ figure captions; the supplied LSUN caption is incomplete.Source: Figure 15 · p. 27, Figure 16 · p. 28, Figure 17 · p. 29

Obs. BAM BAM

• The LSUN inpainting example is identified as the observation and BAM restoration, with the observation displayed as A†y for compressed sensing.The captions also identify BAM⋆ as fine-tuned and use matched yellow boxes for magnification.Source: Obs. BAM BAM · p. 30

E.3 SPARSE-VIEW COMPUTED TOMOGRAPHY

• The experiment uses single-channel 512 × 512 images and a random slice split of 95%/4%/1% for training, validation, and testing.The Gaussian noise level is σy = 10^-4, and the operator uses 51 uniformly spaced projection angles.Source: §E.3 · p. 31
• After CT fine-tuning, a three-step BAM sample recovers fine lung vessels that the RAM estimate smooths out.The experiment uses LIDC-IDRI chest slices with 51 parallel-beam projections.Source: §E.3 · p. 31
• The standard deviation across 64 draws concentrates on anatomical edges and vessels, while the empirical mean reveals hidden structures in the magnified lung region.The draws use a fixed observation; the standard deviation measures variation among generated samples.Source: §E.3 · p. 31
• Multiscale variability and absolute mean error are computed after separately block-averaging each draw and the ground truth at k = 4 and k = 8.The variance is recomputed across downsampled draws rather than by averaging the original standard-deviation map.Source: §E.3 · p. 32, Figure 20 · p. 32

E.4 MOTION DEBLURRING EXPERIMENTS

• For synthetic blind spatially varying blur, BAM and RAM use the same estimated operator, and no ground-truth kernel enters reconstruction.The estimator predicts image-dependent basis kernels and spatial mixing coefficients from the observation; its output remains fixed during reconstruction.Source: §E.4 · p. 33
• Real Köhler observations have no ground truth, so that setting is evaluated qualitatively rather than quantitatively.The motion-deblurring study also includes known spatially invariant blur and synthetic blind spatially varying blur.Source: §E.4 · pp. 32–33

E.5 JPEG RESTORATION

• For noisy JPEG restoration at Q = 10, the observation retains quantization artifacts while network conditioning uses a rounding-relaxed linear approximation.The pre-compression Gaussian noise has standard deviation 0.01; the sampler setting σy = 0.05 is distinct.Source: §E.5 · p. 35

E.6 SINGLE PROBLEM

• For every tested problem, the mean of 64 FFHQ draws has higher PSNR than the average individual draw and exceeds problem-specific RAM⋆ on four of five problems.The evaluation uses σy = 0.025; super-resolution is the exception to the RAM⋆ comparison.Source: §E.6 · p. 36, Table 10 · p. 36

F ROBUSTNESS TO INT8 QUANTIZATION

• The optimization combines INT8 quantization with reformulated physics-aware operations, particularly convolution and transposed-convolution layers in the PhysicsBlock.Profiling identified these operations as a substantial part of inference latency.Source: F ROBUSTNESS TO INT8 QUANTIZATION · p. 37
• 79.91 ms per inference versus 280.56 ms for the baseline gives the INT8-optimized model a 3.51× speedup and 71.5% lower latency.Throughput rises from 3.56 to 12.51 qps, while PSNR changes from 22.49 to 21.20 dB, SSIM from 0.570 to 0.498, and LPIPS from 0.400 to 0.427.Source: F ROBUSTNESS TO INT8 QUANTIZATION · p. 37, Table 11 · p. 37
• The deployment results are obtained without fine-tuning, and the authors leave quantization-aware fine-tuning for future work.The qualitative comparison reports close visual agreement between baseline and INT8-optimized reconstructions.Source: F ROBUSTNESS TO INT8 QUANTIZATION · p. 37, Figure 25 · p. 37