Remove unnecessary deep copies in SBI training/inference workflow

[janfb]

Problem

The SBI codebase currently creates deep copies of neural networks and posteriors in three places:

1. train() returns deepcopy(self._neural_net)
2. build_posterior() returns deepcopy(self._posterior)
3. build_posterior() stores self._model_bank.append(deepcopy(self._posterior))

These deep copies can consume significant memory (10-100MB+ per round for modern density estimators).

Details

• Model bank is completely unused: _model_bank is write-only - no code reads from
• Neural networks aren't reused between rounds: SNPE only needs the data (_theta_roundwise, _x_roundwise) from previous rounds, not the networks
• No modification risk: Users interact with posteriors through stable APIs (sample(), log_prob())

The deep copies appear to be legacy from when multi-round inference was handled internally. Now that users manage rounds explicitly, they serve no purpose.

Proposed Solution

1. Remove _model_bank entirely - it's likely legacy code?
2. Return networks/posteriors without copying:

def train(self, ...):
    return self._neural_net  # No deepcopy

def build_posterior(self, ...):
    # Remove: self._model_bank.append(deepcopy(self._posterior))
    return self._posterior  # No deepcopy

Impact

• Memory savings: ~30x reduction for multi-round inference (e.g., 300MB → 10MB for 10 rounds)
• No functionality loss: All tests should pass unchanged
• Backward compatibility: We could also add return_copy=False parameter to ensure backward compatibility.

Questions

• Any hidden use cases for the model bank we're missing?
• Should we make copying opt-in via parameter or just remove it?

[michaeldeistler]

Amazing! Love this! And yes, I think most of this is either legacy or based on decisions that were not all that much thought through. The only thing I am not sure about is:

SNPE only needs the data (_theta_roundwise, _x_roundwise) from previous rounds, not the networks.

I think SNPE-B needs the network from previous rounds for proposal log-probs, no?

[janfb]

Yes, on a closer look, SNPE-C needs the snapshots of the neural net weights from previous rounds when evaluating the mog-based proposal in non-atomic SNPE-C, e.g., here:

sbi/sbi/inference/trainers/npe/npe_c.py

Line 178 in ebcd68e

proposal = self._proposal_roundwise[-1]

and then

sbi/sbi/inference/trainers/npe/npe_c.py

Lines 432 to 435 in ebcd68e

	encoded_x = proposal.posterior_estimator.net._embedding_net(proposal.default_x)
	dist = (
	proposal.posterior_estimator.net._distribution
	) # defined to avoid ugly black formatting.

And SNPE-B needs to evaluate the log prob of the proposals from previous rounds for calculating the importance weights:

sbi/sbi/inference/trainers/npe/npe_b.py

Lines 114 to 120 in ebcd68e

	for k, density in enumerate(self._proposal_roundwise):
	# we accept the k th proposal log prob to be -Inf at theta
	# meaning that theta is out of the k th proposal range
	log_previous_proposals[:, k] = density.log_prob(theta)
	utils.assert_not_nan_or_plus_inf(
	log_previous_proposals[:, k], "proposal log probs of proposal samples"
	)

Then question is where we create the deepcopy, in train() or in build_posterior()?

We can definitively remove the _model_bank. Then we could return a deepcopy in train() and return the plain posterior with the deep copied net attached in build_posterior?