From 1a07bc7eed236cf150a271771bd6897f38782f2d Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Tue, 14 Jul 2026 08:39:48 -0700 Subject: [PATCH] Two independent memory reductions for SWIFT on high-dimensional data. 1. Memory-efficient marginal oracle for final estimation. SWIFT's final MirrorDescent estimation hardcoded the HUGIN marginal oracle (message_passing_stable), which materializes every junction-tree clique belief simultaneously. On high-dimensional data with large maximal cliques this drives peak memory to the sum of all clique beliefs (tens of GiB) and is numerically unstable with -inf potentials (HUGIN uses belief subtraction). Switch to message_passing_implicit with the einsum_materialized contraction: - Peak memory is bounded by the largest single clique super-factor rather than the sum of all beliefs. - Numerically stable: log-space add + logsumexp, no exp/log round-trip and no belief subtraction. - Fastest contraction on GPU. Both oracles perform exact junction-tree inference, so estimated marginals are identical up to floating-point error. 2. Free precomputed candidate marginals after measurement. The candidate marginals produced by from_projectable (plus the workload intermediates) are only needed through query selection and measurement, but they stayed live for the rest of __call__ -- through estimation and the column-by-column generation phase. On host-memory-constrained slices this residual is enough to push generation over the host RAM limit. Free them right after measurement so the peak during estimation and generation is lower. PiperOrigin-RevId: 947706110 --- dpsynth/discrete_mechanisms/swift.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/dpsynth/discrete_mechanisms/swift.py b/dpsynth/discrete_mechanisms/swift.py index 62bcc8a..73a1f34 100644 --- a/dpsynth/discrete_mechanisms/swift.py +++ b/dpsynth/discrete_mechanisms/swift.py @@ -197,8 +197,14 @@ def __call__( ######################################################## # Precompile MirrorDescent + synth while measuring. # ######################################################## + # Use implicit-factor message passing so peak memory is bounded by the + # largest single clique super-factor rather than the sum of all clique + # beliefs (as in the HUGIN oracle). einsum_materialized is numerically + # stable and the fastest contraction on GPU. closed_oracle = functools.partial( - mbi.marginal_oracles.message_passing_stable, jtree=jtree + mbi.marginal_oracles.message_passing_implicit, + jtree=jtree, + contraction=mbi.marginal_oracles.einsum_materialized, ) estimator = mbi.estimation.MirrorDescent(marginal_oracle=closed_oracle) rows = int(mbi.estimation.minimum_variance_unbiased_total(measurements)) @@ -224,6 +230,15 @@ def __call__( measurements.extend(new_measurements) logging.info('[SWIFT] Finished measurements.') + # The precomputed candidate marginals (`answers`) and the workload + # intermediates are dead once measurement is done, but they hold a large + # amount of host RAM (all candidate marginals materialized by + # `from_projectable`). Estimation and, especially, column-by-column + # generation run afterwards; on host-memory-constrained slices the residual + # `answers` is enough to push generation over the host RAM limit. Free them + # now so the peak during estimation/generation is lower. + del answers, errors, candidates + ######################################################## # Estimate the model using all measurements # ########################################################