Predict x0 from a Noise Estimate

Problem Statement

Every sampler step starts by reconstructing the clean sample from the current x_t and the network's noise prediction. Implement that estimate.

Background

Inverting the forward equation for x0 given the predicted noise eps:

$\hat{x}_0 = \frac{x_t - \sqrt{1 - \bar{\alpha}_t}\, \hat{\varepsilon}}{\sqrt{\bar{\alpha}_t}}$

Your Task

Implement:

def predict_x0(x_t, eps, alpha_bar_t):

Return x0_hat as a list rounded to 4 decimals.

Input Format

• x_t, eps: lists of length D.
• alpha_bar_t (float) in (0, 1].

Output Format

• A list of D floats rounded to 4 decimals.

Sample

print(predict_x0([1.4142, 0.0], [1.0, -1.0], 0.5))

Output:

[1.0, 1.0]