SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(2) In addition, we ground the performance of a sampler in the downstream task of classifying a partially observed image  $x_{\mathrm{M}}^{0} = y$ : we use a classifier to predict the class  $\hat{z}$  of an inpainted image  $x^{0}$ , and compute the accuracy against the true class  $z^{*}$  of the unmasked image ( $z^{*}$  is provided in MNIST dataset). This prediction is made by the same classifier used in Section 5.2.
Consider  $K$  weighted particles  $\{x_k^0; w_k^0\}_{k=1}^K$  drawn from a sampler conditioned on  $x_{\mathbf{M}}^0 = y$  and assume the weights are normalized. We define the Bayes accuracy (BA) as
$$
\mathbb {1} \left\{\hat {z} (y) = z ^ {*} \right\}, \quad \text {w i t h} \quad \hat {z} (y) := \underset {z = 1: 1 0} {\arg \max } \sum_ {k = 1} ^ {K} w _ {k} ^ {0} p (z; x _ {k} ^ {0}),
$$
where  $\hat{z}(y)$  is viewed as an approximation to the Bayes optimal classifier  $\hat{z}^*(y)$  given by
$$
\hat {z} ^ {*} (y) := \underset {z = 1: 1 0} {\arg \max } p _ {\theta} (z \mid x _ {\mathbf {M}} ^ {0} = y) = \underset {z = 1: 1 0} {\arg \max } \int p _ {\theta} \left(x ^ {0} \mid x _ {\mathbf {M}} ^ {0} = y\right) p \left(z; x ^ {0}\right) d x ^ {0}. \tag {18}
$$
(In eq. (18) we assume the classifier  $p_{y|x^0}(\cdot |x^0)$  is the optimal classifier on full images.)
(3) We also consider classification accuracy (CA) defined as the following
$$
\sum_ {k = 1} ^ {K} w _ {k} ^ {0} \mathbb {1} \left\{\hat {z} \left(x _ {k} ^ {0}\right) = z ^ {*} \right\}, \quad \text {w i t h} \quad \hat {z} \left(x _ {k} ^ {0}\right) := \underset {z = 1: 1 0} {\arg \max } p _ {y | x ^ {0}} \left(z | x _ {k} ^ {0}\right).
$$
BA and CA evaluate different aspects of a sampler. BA is focused on the optimal prediction among multiple particles, whereas CA is focused on the their weighted average prediction.
Comparison results of different methods. Figure I depicts the ESS trace, BA and CA for different samplers. The overall observations are similar to the observations in the class-conditional generation task in Section 5.2, except that SMC-Diff and Replacement methods are not available there.
Replacement has the lowest CA and BA across all settings. Comparing TDS to SMC-Diff, we find that SMC-Diff's ESS is consistently greater than TDS; however, SMC-Diff is outperformed by TDS in terms of both CA and BA.
We also note that despite Gradient Guidance's CA is lower, its BA is comparable to TDS. This result is due to that as long as Gradient Guidance generates a few good-quality samples out of  $K$  particles, the optimal prediction can be accurate, thereby resulting in a high BA.
Ablation study on twisting function variance schedule. We compare the following three options:
1. DPS:  $\hat{\sigma}_t\coloneqq 2\| \hat{x}_\theta (x^t)_{\mathbf{M}} - y\| _2\sigma_t^2$  , adapted from the DPS method [6, Algorithm 1],
2.  $\Pi GDM$ :  $\hat{\sigma}_t \coloneqq \sigma_t^2 / \sqrt{\bar{\alpha}_t}$ , adapted from the  $\Pi GDM$  method [32, Algorithm 1],
3. TDS:  $\hat{\sigma}_t^2 = \frac{(\bar{\sigma}_t^2 / \bar{\alpha}_t)\cdot\tau^2}{(\bar{\sigma}_t^2 / \bar{\alpha}_t) + \tau^2}$  where  $\tau^2 = 0.12$
Figure J shows the classification accuracy of TDS, TDS-IS and Gradient Guidance with different twist scale schemes. We find that our choice has similar performance to that of IIGDM, and outperforms DPS in most cases. Exceptions are when  $\mathbf{M} =$  "quarter" and for large  $K$ , TDS with twist scale choice of IIGDM or DPS has higher CA, as is shown in the left panel in Figure Jb.
D.2.3 Class-conditional generation on ImageNet and CIFAR-10
ImageNet. We compare TDS (with # of particles  $K = 1, 16$ ) to Classifier Guidance (CG) using the same unconditional model from https://github.com/openai/guided-diffusion/tree/main.CG uses a classifier trained on noisy inputs taken from this repository as well. For TDS, we
(a)  $\mathbf{M} =$  "half".ESStrace.
(c)  $\mathbf{M} =$  "quarter". ESS trace.
(b)  $\mathbf{M} =$  "half". Class. accuracy (left), and Bayes accuracy (right) v.s.  $K$ .
(d)  $\mathbf{M} =$  "quarter". Class. accuracy (left), and Bayes accuracy (right) v.s.  $K$ .
Figure I: MNIST image inpainting. Results for observed dimension  $\mathbf{M} =$  "half" are shown in the top panel, and  $\mathbf{M} =$  "quarter" in the bottom panel. (i) Left column: ESS traces are averaged over 100 different images inpainted by TDS, TDS-IS, and SMC-Diff, all with  $K = 64$  particles. TDS's ESS is generally increasing until the final 20 steps where it drops to around 1, suggesting significant particle collapse in the end of generation trajectory. TDS-IS's ESS is always around 1. SMC-Diff has higher particle efficiency compared to TDS. (ii) Right column: Classification accuracy and Bayes accuracy are averaged over 10,000 images. In general increasing the number of particles  $K$  would improve the performance of all samplers. TDS and TDS-truncate have the highest accuracy among all given the same  $K$ . (iii) Finally, comparison of top and bottom panels shows that in a harder inpainting problem where  $\mathbf{M} =$  "quarter", TDS's has a higher ESS but lower CA and BA.
Method	Classification Accuracy ↑	FID ↓	Inception Score ↑
TDS (K=16)	99.90%	26.65	64.03
TDS (K=1)	99.63%	26.05	44.45
Classifier Guidance	99.17%	14.03	100.33
Unconditional model	n/a	26.21	39.70
Table 1: ImageNet. Comparison of sample quality. The top three methods are evaluated on 16k samples generated with 100 steps. The unconditional model performance is provided in [9], which is evaluated on 50k samples generated with 250 steps.
use the same classifier evaluated at timestep  $= 0$  to mimic a standard trained classifier. We generate 16 images for each of the 1000 class labels, using 100 sampling steps. TDS uses a twist scale of 10, and CG uses a guidance scale of 10. Notably, given a fixed class, TDS ( $K = 16$ ) generates correlated samples in a single SMC run, and TDS ( $K = 1$ ) and CG generate 16 independent samples.
In Figure K, we observe that TDS can faithfully capture the class and have comparable image quality to CG's, although with less diversity than CG and TDS  $(K = 1)$ .
Figure L shows more samples given randomly selected classes. We also reported results of the unconditional model from the original paper [9] that are evaluated on 50k samples with 250 sampling steps in Appendix D.2.3. TDS and CG provide similar classification accuracy. TDS has similar FIDs compared to the unconditional model and better inception score. CG's FID and inception score are better than TDS. We suspect this difference is attributed to the sample correlation (and hence less diversity) within particles in a single run of TDS ( $K = 16$ ).
CIFAR-10. We ran TDS  $(K = 16)$  with the twist scale  $= 1$  and 100 sampling steps using diffusion model from https://github.com/openai/improved-diffusion and classifier from https: //github.com/VSehwag/minimal-diffusion/tree/main. TDS generates faithful and diverse
(a)  $\mathbf{M} =$  "half". Classification accuracy under different twist scale schemes.
(b)  $\mathbf{M} =$  "quarter". Classification accuracy under different twist scale schemes.
Figure J: MNIST image inpainting: Ablation study on twist scales. Top: observed dimensions  $\mathbf{M} =$  "half". Bottom:  $\mathbf{M} =$  "quarter". In most case, the performance of our choice of twist scale is similar to that of IIGDM, and is better compared to DPS.
(a) TDS  $(\mathrm{P} = 16)$
(b) TDS  $(\mathrm{P} = 1)$
(c) Classifier Guidance
Figure K:  $256 \times 256$  ImageNet. Samples from TDS and Classifier Guidance given class 'brambling'.
images, see Figure M. However, we found TDS can occasionally generate incorrect samples for the class 'truck' (the last panel in Figure M).
Finally, we present the ESS trace plot of the three image datasets in Figure N. We see there is a sudden ESS drop in the final stage for MINST, sometimes for CIFAR10, but usually not on ImageNet. Mechanically, the drop in ESS implies a large discrepancy between the final and intermediate target distributions. We suspect such discrepancies might arise from irregularities in the denoising network near  $t = 0$ . In practice, one can consider early truncating the TDS procedure to prevent such ESS drops and promote particle diversity, as is studied in Appendices D.2.1 and D.2.2.
E Motif-scaffolding application details
Unconditional model of protein backbones. We here use the FrameDiff, a diffusion generative model described by Yim et al. [39]. FrameDiff parameterizes protein  $N$ -residue protein backbones as a collection of rigid bodies defined by rotations and translations as  $x^0 = [(R_1,z_1),\ldots ,(R_N,z_N)]\in$
Figure L:  $256 \times 256$  ImageNet. 15 random class samples. Top to bottom: TDS(P=16), TDS(P=1), Classifier Guidance.
(a) airplane
(b) deer
(c) horse
(d) ship
(e) truck
(f) failure mode
Figure M: CIFAR10. Samples from TDS with  $K = 16$  particles given different class labels (within a single SMC run). In most cases TDS generate authentic and diverse images. The last panel presents a failure mode for the class 'truck'.
Figure N: 2 example effective sample size traces of TDS ( $K = 16$ , 100 sampling steps) on MNIST, CIFAR-10 and ImageNet models, respectively.
$SE(3)^N$ .  $SE(3)$  is the special Euclidean group in three dimensions (a Riemannian manifold). Each  $R_{n} \in SO(3)$  (the special orthogonal group in three dimensions) is a rotation matrix and each  $z_{n} \in \mathbb{R}^{3}$  in a translation. Together,  $R_{n}$  and  $z_{n}$  describe how one obtains the coordinates of the three backbone atoms  $\mathbb{C}, \mathbb{C}_{\alpha}$ , and  $\mathbb{N}$  for each residue by translating and rotating the coordinates of an idealized residue with  $C_{\alpha}$  carbon and the origin. The conditioning information is then a motif  $y = x_{\mathbf{M}}^{0} \in SE(3)^{|\mathbf{M}|}$  for some  $\mathbf{M} \subset \{1, \ldots, N\}$ . FrameDiff is a continuous time diffusion model and includes the number of steps as a hyperparameter; we use 200 steps in all experiments. We refer the reader to [39] for details on the neural network architecture, and details of the forward and reverse diffusion processes.
Twisting functions for motif scaffolding. To define a twisting function, we use a tangent normal approximation to  $p_{\theta}(y \mid x^t)$  as introduced in eq. (15). In this case the tangent normal factorizes across each residue and across the rotational and translational components. The translational components are represented in  $\mathbb{R}^3$ , and so are treated as in the previous Euclidean cases, using the variance preserving extension described in Appendix A. For the rotational components, represented in the special orthogonal group  $SO(3)$ , the tangent normal may be computed as described in Appendix B, using the known exponential map on  $SO(3)$ . In particular, we use as the twisting function
$$
\tilde {p} _ {\boldsymbol {\theta}} (y \mid x ^ {t}, \mathbf {M}) = \prod_ {m = 1} ^ {| \mathbf {M} |} \mathcal {N} (y _ {m} ^ {T}; \hat {x} _ {\boldsymbol {\theta}} (x ^ {t}) _ {\mathbf {M} _ {m}} ^ {T}, 1 - \bar {\alpha} _ {t}) \mathcal {T N} _ {\hat {x} _ {\boldsymbol {\theta}} (x ^ {t}) _ {\mathbf {M} _ {m}} ^ {R}} (y _ {m} ^ {R}; 0, \bar {\sigma} _ {t} ^ {2}),
$$
where  $y_{m}^{T}, \hat{x}_{\theta}(x^{t})_{\mathbf{M}_{m}}^{T} \in \mathbb{R}^{3}$  represent the translations associated with the  $m^{th}$  residue of the motif and its prediction from  $x^{t}$ , and  $y_{m}^{R}, \hat{x}_{\theta}(x^{t})_{\mathbf{M}_{m}}^{R} \in SO(3)$  represent the analogous quantities for the rotational component. Next,  $1 - \bar{\alpha}_{t} = \mathrm{Var}_{q}[x^{t} \mid x^{0}]$  and  $\bar{\sigma}_{t}^{2}$  is the time of the Brownian motion associated with the forward process at step  $t$  [39]. For further simplicity, our implementation further approximates log density of the tangent normal on rotations using the squared Frobenius norm of the difference between the rotations associated with the motif and the denoising prediction (as in [37]), which becomes exact in the limit that  $\bar{\sigma}_{t}$  approaches 0 but avoids computation of the inverse exponential map and its Jacobian.
Motif rotation and translation degrees of freedom. We similarly seek to eliminate the rotation and translation of the motif as a degree of freedom. We again represent our ambivalence about the rotation of the motif with randomness, augment our joint model to include the rotation with a uniform prior, and write
$$
p _ {y \mid x ^ {0}} (y \mid x ^ {0}) = \int p (R) p _ {y \mid x ^ {0}} (y \mid x ^ {0}, R) \quad \text {f o r} \quad p _ {y \mid x ^ {0}} (y \mid x ^ {0}, R) = \delta_ {R y} (x ^ {0})
$$
and with  $p(R)$  the uniform density (Haar measure) on SO(3), and  $Ry$  represents rotating the rigid bodies described by  $y$  by  $R$ . For computational tractability, we approximate this integral with a Monte Carlo approximation defined by subsampling a finite number of rotations,  $\mathcal{R} = \{R_k\}$  (with  $|\mathcal{R}| = \#$  Motif Rots. in total). Altogether we have
$$
\tilde {p} _ {\theta} (y \mid x ^ {t}) = | \mathcal {R} | ^ {- 1} \cdot | \mathcal {M} | ^ {- 1} \sum_ {R \in \mathcal {R}} \sum_ {\mathbf {M} \in \mathcal {M}} \tilde {p} _ {\theta} (R y \mid x ^ {t}, \mathbf {M}).
$$
The subsampling above introduces the number of motif locations and rotations subsampled as a hyperparameter. Notably the conditional training approach does not immediately address these degrees of freedom, and so prior work randomly sampled a single value for these variables [34, 36, 37].
Evaluation details. We use self-consistency evaluation approach for generated backbones [34] that (i) uses fixed backbone sequence design (inverse folding) to generate a putative amino acid sequence to encode the backbone, (ii) forward folds sequences to obtain backbone structure predictions, and (iii) judges the quality of initial backbones by their agreement (or self-consistency) with predicted structures. We inherit the specifics of our evaluation and success criteria set-up following [37], including using ProteinMPNN [7] for step (i) and AlphaFold [20] on a single sequence (no multiple sequence alignment) for (ii).
In this evaluation we use ProteinMPNN [7] with default settings to generate 8 sequences for each sampled backbone. Positions not indicated as resdesignable in [37, Methods Table 9] are held fixed. We use AlphaFold [20] for forward folding. We define a "success" as a generated backbone for which at least one of the 8 sequences has backbone atom with both scRMSD  $< 1\AA$  on the motif and scRMSD  $< 2\AA$  on the full backbone.
We benchmarked TDS on 24/25 problems in the benchmark set introduced [37, Methods Table 9]. A 25th problem (6VW1) is excluded because it involves multiple chains, which cannot be represented by FrameDiff. FrameDiff requires specifying a total length of scaffolds. In all replicates, we fixed the scaffold the median of the Total Length range specified by Watson et al. [37, Methods Table 9]. For example, 116 becomes 125 and 62-83 becomes 75.
As discussed in the main text, in our handling of the motif-placement degree of freedom we restrict the number of possible placements considered for discontinuous motifs. To select these placements we (i) restrict to masks which place the motif indices in a pre-specified order that we do not permute and do not separate residues that appear contiguously in source PDB file, and (ii) when there are still too many possible placements sub-sample randomly to obtain the set of masks  $\mathcal{M}$  of at most some maximum size (# Motif Locs.). We do not enforce the spacing between segments specified in the "contig" description described by [37].
Resampling. In all motif-scaffolding experiments, we use systematic resampling. An effective sample size threshold is used  $K / 2$  in all cases, that is, we trigger resampling steps only when the effective sample size falls below this level.
Figure O: Results on full motif-scaffolding benchmark set. The y-axis is the fraction of successes across at least 200 replicates. Error bars are  $\pm 2$  standard errors of the mean. Problems with scaffolds shorter than 100 residues are bolded. TDS  $(\mathrm{K} = 8)$  outperforms the state-of-the-art (RFdiffusion) on most problems with short scaffolds.
Figure P: Increasing the twist scale has a different impact across motif-scaffolding benchmark problems.
E.1 Additional Results
Impact of twist scale on additional motifs. Figure 3a showed monotonically increasing success rates with the twist scale. However, this trend does not hold for every problem. Figure P demonstrates this by comparing the success rates with different twist-scales on five additional benchmark problems.
Effective sample size varies by problem. Figure Q shows two example effective sample size traces over the course of sample generation. For 6EXZ-med resampling was triggered 38 times (with 14 in the final 25 steps), and for 5UIS resampling was triggered 63 times (with 13 in the final 25 steps). The traces are representative of the larger benchmark.
Application of TDS to RFdiffusion: We also tried applying TDS to RFdiffusion [37]. RFdiffusion is a diffusion model that uses the same backbone structure representation as FrameDiff. However, unlike FrameDiff, RFdiffusion trained with a mixture of conditional examples, in which a segment of the backbone is presented as input as a desired motif, and unconditional training examples in which the only a noised backbone is provided. Though in our benchmark evaluation provided the motif information as explicit conditioning inputs, we reasoned that TDS should apply to RFdiffusion as well (with conditioning information not provided as input). However, we were unable to compute numerically stable gradients (with respect to either the rotational or translational components of the backbone representation); this lead to twisted proposal distributions that were similarly unstable and trajectories that frequently diverged even with one particle. We suspect this instability owes to RFdiffusion's structure prediction pretraining and limited fine-tuning, which may allow it to achieve good performance without having fit a smooth score-approximation.
Figure Q: Effective sample size traces for two motif-scaffolding examples (Left) 6EXZ-med and (Right) 5IUS. In both case  $K = 8$  and a resampling threshold of  $0.5K$  is used. Dashed red lines indicate resampling times.