SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We refer to Yellott (1977) or Grover et al. (2019) for the proof. However, Grover et al. (2019) provide only an adapted proof sketch from Yellott (1977). The probability of sampling element  $i$  first is given by its score  $s_i$  divided by the sum of all weights in the set
$$
q \left(\tilde {s} _ {i}\right) = \frac {s _ {i}}{Z} \tag {15}
$$
For  $z_{i} = \log s_{i}$ , the right hand side of Equation (15) is equal to the softmax distribution  $\mathrm{softmax}(z_{i}) = \exp(z_{i}) / \sum_{j} \exp(z_{j})$  as already described in (Xie and Ermon, 2019). Hence, Equation (15) directly leads to the Gumbel-Softmax trick (Jang et al., 2016; Maddison et al., 2017).
A.2.1 Differentiable Sorting
In the main text of the paper we rely on a differentiable function  $f_{\pi}(\tilde{s})$ , which sorts the resampled version of the scores  $s$
$$
\pi = f _ {\pi} (\tilde {s}) = \operatorname {s o r t} (\tilde {s}) \tag {16}
$$
Here, we summarise the findings from Grover et al. (2019) on how to construct such a differentiable sorting operator. As already mentioned in Section 2, there are multiple works on the topic (Prillo and Eisenschlos, 2020; Petersen et al., 2021; Mena et al., 2018), but we restrict ourselves to the work of Grover et al. (2019) as we see the differentiable generation of permutation matrices as a tool in our pipeline.
Corollary A.2 (Permutation Matrix (Grover et al., 2019)). Let  $\boldsymbol{s} = [s_1, \dots, s_n]^T$  be a real-valued vector of length  $n$ . Let  $A_s$  denote the matrix of absolute pairwise differences of the elements of  $\boldsymbol{s}$  such that  $A_s[i, j] = |s_i - s_j|$ . The permutation matrix  $\pi$  corresponding to  $\text{sort}(s)$  is given by:
$$
\pi = \left\{ \begin{array}{l l} 1 & \text {i f} j = \arg \max  [ (n + 1 - 2 i) s - A _ {s} \mathbb {1} ] \\ 0 & \text {o t h e r w i s e} \end{array} \right. \tag {17}
$$
where  $\mathbb{1}$  denotes the column vector of all ones.
As we know, the arg max operator is non-differentiable which prohibits the direct use of Corollary A.2 for gradient computation. Hence, Grover et al. (2019) propose to replace the arg max operator with softmax to obtain a continuous relaxation  $\pi(\tau)$  similar to the GS trick (Jang et al., 2016; Maddison et al., 2017). In particular, the  $i$ th row of  $\pi(\tau)$  is given by:
$$
\pi (\tau) [ i,: ] = \operatorname {s o f t m a x} [ (n + 1 - 2 i) s - A _ {s} \mathbb {1} / \tau ] \tag {18}
$$
where  $\tau > 0$  is a temperature parameter. We adapted this section from Grover et al. (2019) and we also refer to their original work for more details on how to generate differentiable permutation matrices.
In this, work we remove the temperature parameter  $\tau$  to reduce clutter in the notation. Hence, we only write  $\pi$  instead of  $\pi(\tau)$ , although it is still needed for the generation of the matrix  $\pi$ . For details on how we select the temperature parameter  $\tau$  in our experiments, we refer to Appendix C.
B Detailed Derivation of the Differentiable Two-Stage Random Partition Model
B.1 Two-Stage Partition Model
We want to partition  $n$  elements  $[n] = \{1, \dots, n\}$  into  $K$  subsets  $\{\mathcal{S}_1, \dots, \mathcal{S}_K\}$  where  $K$  is a priori unknown.
Definition B.1 (Partition). A partition  $\rho$  of a set of elements  $[n] = \{1, \ldots, n\}$  is a collection of subsets  $(S_1, \ldots, S_K)$  such that
$$
\mathcal {S} _ {1} \cup \dots \cup \mathcal {S} _ {K} = [ n ] \quad a n d \quad \forall i \neq j: \mathcal {S} _ {i} \cap \mathcal {S} _ {j} = \emptyset \tag {19}
$$
Put differently, every element  $i$  has to be assigned to precisely one subset  $\mathcal{S}_k$ . We denote the size of the  $k$ -th subset  $\mathcal{S}_k$  as  $n_k = |\mathcal{S}_k|$ . Alternatively, we describe a partition  $\rho$  as an assignment matrix  $Y = [\pmb{y}_1,\dots,\pmb{y}_K]^T\in \{0,1\}^{K\times n}$ . Every row  $\pmb{y}_k\in \{0,1\}^{1\times n}$  is a multi-hot vector, where  $\pmb{y}_{ki} = 1$  assigns element  $i$  to subset  $\mathcal{S}_k$ .
In this work, we propose a new two-stage procedure to learn partitions. The proposed formulation separately infers the number of elements per subset  $n_k$  and the assignment of elements to subsets  $S_k$  by inducing an order on the  $n$  elements and filling  $S_1, \ldots, S_K$  sequentially in this order. See Figure 1 for an example.
Definition B.2 (Two-stage partition model). Let  $\mathbf{n} = [n_1, \ldots, n_K] \in \mathbb{N}_0^K$  be the subset sizes in  $\rho$ , with  $\mathbb{N}_0$  the set of natural numbers including 0 and  $\sum_{k=1}^{K} n_k = n$ , where  $n$  is the total number of elements. Let  $\pi \in \{0,1\}^{n \times n}$  be a permutation matrix that defines an order over the  $n$  elements. We define the two-stage partition model of  $n$  elements into  $K$  subsets as an assignment matrix  $Y = [\mathbf{y}_1, \ldots, \mathbf{y}_K]^T \in \{0,1\}^{K \times n}$  with
$$
\boldsymbol {y} _ {k} = \sum_ {i = \nu_ {k} + 1} ^ {\nu_ {k} + n _ {k}} \boldsymbol {\pi} _ {i}, \quad \text {w h e r e} \quad \nu_ {k} = \sum_ {\iota = 1} ^ {k - 1} n _ {\iota} \tag {20}
$$
such that  $Y = \left[\left\{\boldsymbol{y}_k \mid n_k > 0\right\}_{k=1}^K\right]^T$ .
Note that in contrast to previous work on partition models (Mansour and Schork, 2016), we allow  $S_{k}$  to be the empty set  $\emptyset$ . Hence,  $K$  defines the maximum number of possible subsets, not the effective number of non-empty subsets.
To model the order of the elements, we use a permutation matrix  $\pi = [\pi_1,\dots ,\pi_n]^T\in \{0,1\}^{n\times n}$  which is a square matrix where every row and column sums to 1. This doubly-stochastic property of all permutation matrices  $\pi$  (Marcus, 1960) thus ensures that the columns of  $Y$  remain one-hot vectors. At the same time, its rows correspond to  $n_k$ -hot vectors  $\pmb{y}_k$  in Definition B.2 and therefore serve as subset assignment vectors.
Corollary B.3. A two-stage partition model  $Y$ , which follows Definition B.2, is a valid partition satisfying Definition B.1.
Proof. By definition, every row  $\pi_{i}$  and column  $\pi_{j}$  of  $\pi$  is a one-hot vector, hence every  $\sum_{i = \nu_k + 1}^{\nu_k + n_k}\pi_i$  results in different, non-overlapping  $n_k$ -hot encodings, ensuring  $\mathcal{S}_i\cap \mathcal{S}_j = \emptyset \forall i,j$  and  $i\neq j$ . Further, since  $n_k$ -hot encodings have exactly  $n_k$  entries with 1, we have  $\sum_{i = \nu_k + 1}^{\nu_k + n_k}\sum_{j = 1}^{n}\pi_{ij} = n_k$ . Hence, since  $\sum_{k = 1}^{K}n_{k} = n$ , every element  $i$  is assigned to a  $\pmb{y}_k$ , ensuring  $\mathcal{S}_1\cup \dots \cup \mathcal{S}_K = [n]$ .
B.2 Two-Stage Random Partition Models
An RPM  $p(Y)$  defines a probability distribution over partitions  $Y$ . In this section, we derive how to extend the two-stage procedure from Definition B.2 to the probabilistic setting to create a two-stage RPM. To derive the two-stage RPM's probability distribution  $p(Y)$ , we need to model distributions over  $n$  and  $\pi$ . We choose the MVHG distribution  $p(n; \omega)$  and the PL distribution  $p(\pi; s)$  (see Section 3).
We calculate the probability  $p(Y; \boldsymbol{\omega}, s)$  sequentially over the probabilities of subsets  $p_{\boldsymbol{y}_k} := p(\boldsymbol{y}_k \mid \boldsymbol{y}_{<k}; \boldsymbol{\omega}, s)$ .  $p_{\boldsymbol{y}_k}$  itself depends on the probability over subset permutations  $p_{\bar{\pi}_k} := p(\bar{\pi} \mid n_k, \boldsymbol{y}_{<k}; s)$ , where a subset permutation matrix  $\bar{\pi}$  represents an ordering over  $n_k$  out of  $n$  elements.
Definition B.4 (Subset permutation matrix  $\bar{\pi}$ ). A subset permutation matrix  $\bar{\pi} \in \{0,1\}^{n_k \times n}$ , where  $n_k \leq n$ , must fulfill
$$
\forall i \leq n _ {k}: \sum_ {j = 1} ^ {n} \bar {\pi} _ {i j} = 1 a n d \forall j \leq n: \sum_ {i = 1} ^ {n _ {k}} \bar {\pi} _ {i j} \leq 1.
$$
We describe the probability distribution over subset permutation matrices  $p_{\bar{\pi}_k}$  using Definition B.4 and Equation (3).
Lemma B.5 (Probability over subset permutations  $p_{\bar{\pi}_k}$ ). The probability  $p_{\bar{\pi}_k}$  of any subset permutation matrix  $\bar{\pi} = [\bar{\pi}_1, \dots, \bar{\pi}_{n_k}]^T \in \{0, 1\}^{n_k \times n}$  is given by
$$
p _ {\bar {\pi} _ {k}} := p \left(\bar {\pi} \mid n _ {k}, \boldsymbol {y} _ {<   k}; \boldsymbol {s}\right) = \prod_ {i = 1} ^ {n _ {k}} \frac {\left(\bar {\pi} \boldsymbol {s}\right) _ {i}}{Z _ {k} - \sum_ {j = 1} ^ {i - 1} \left(\bar {\pi} \boldsymbol {s}\right) _ {j}} \tag {21}
$$
where  $\pmb{y}_{< k} = \{\pmb{y}_1, \dots, \pmb{y}_{k-1}\}$ ,  $Z_k = Z - \sum_{j \in S_{< k}} \pmb{s}_j$  and  $S_{< k} = \bigcup_{j=1}^{k-1} S_j$ .
Proof. We provide the proof for  $p_{\bar{\pi}_1}$ , but it is equivalent for all other subsets. Without loss of generality, we assume that there are  $n_1$  elements in  $S_1$ . Following Equation (3), the probability of a permutation matrix  $p(\pi; s)$  is given by
$$
p (\pi ; \boldsymbol {s}) = \frac {(\pi \boldsymbol {s}) _ {1}}{Z} \frac {(\pi \boldsymbol {s}) _ {2}}{Z - (\pi \boldsymbol {s}) _ {1}} \dots \frac {(\pi \boldsymbol {s}) _ {n}}{Z - \sum_ {j = 1} ^ {n - 1} (\pi \boldsymbol {s}) _ {j}} \tag {22}
$$
At the moment, we are only interested in the ordering of the first  $n_1$  elements. The probability of the first  $n_1$  is given by marginalizing over the remaining  $n - n_1$  elements:
$$
p (\bar {\pi} \mid n _ {1}; \boldsymbol {\omega}) = \sum_ {\pi \in \Pi_ {1}} p (\pi \mid s) \tag {23}
$$
where  $\Pi_1$  is the set of permutation matrices such that the top  $n_1$  rows select the elements in a specific ordering  $\bar{\pi} \in \{0, 1\}^{n_1 \times n}$ , i.e.  $\Pi_1 = \{\pi : [\pi_1, \dots, \pi_{n_1}]^T = \bar{\pi}\}$ . It follows
$$
\begin{array}{l} p (\bar {\pi} \mid n _ {1}; \boldsymbol {\omega}) = \sum_ {\pi \in \Pi_ {1}} p (\pi \mid s) (24) \\ = \sum_ {\pi \in \Pi_ {1}} \prod_ {i = 1} ^ {n} \frac {(\pi \boldsymbol {s}) _ {i}}{Z - \sum_ {j = 1} ^ {i - 1} (\pi \boldsymbol {s}) _ {j}} (25) \\ = \prod_ {i = 1} ^ {n _ {1}} \frac {(\bar {\pi} \boldsymbol {s}) _ {i}}{Z - \sum_ {j = 1} ^ {i - 1} (\bar {\pi} \boldsymbol {s}) _ {j}} \sum_ {\pi \in \Pi_ {1}} \prod_ {i = 1} ^ {n - n _ {1}} \frac {(\pi \boldsymbol {s}) _ {n _ {1} + i}}{Z - \sum_ {j = 1} ^ {n _ {1}} (\bar {\pi} \boldsymbol {s}) _ {j} - \sum_ {j = 1} ^ {i - 1} (\bar {\pi} \boldsymbol {s}) _ {j}} (26) \\ = \prod_ {i = 1} ^ {n _ {1}} \frac {(\bar {\pi} \boldsymbol {s}) _ {i}}{Z - \sum_ {j = 1} ^ {i - 1} (\bar {\pi} \boldsymbol {s}) _ {j}} \sum_ {\pi \in \Pi_ {1}} \prod_ {i = 1} ^ {n - n _ {1}} \frac {(\pi \boldsymbol {s}) _ {n _ {1} + i}}{Z _ {1} - \sum_ {j = 1} ^ {i - 1} (\bar {\pi} \boldsymbol {s}) _ {j}} (27) \\ \end{array}
$$
where  $Z_{1} = Z - \sum_{j = 1}^{n_{1}}(\bar{\pi}\pmb {s})_{j}$  . It follows
$$
p (\bar {\pi} \mid n _ {1}; \boldsymbol {\omega}) = \prod_ {i = 1} ^ {n _ {1}} \frac {(\bar {\pi} \boldsymbol {s}) _ {i}}{Z - \sum_ {j = 1} ^ {i - 1} (\bar {\pi} \boldsymbol {s}) _ {j}} \tag {28}
$$
Lemma B.5 describes the probability of drawing the elements  $i \in S_k$  in the order described by the subset permutation matrix  $\bar{\pi}$  given that the elements in  $S_{<k}$  are already determined. Note that in a slight abuse of notation, we use  $p(\bar{\pi} \mid n_k, \boldsymbol{y}_{<k}; \boldsymbol{\omega}, \boldsymbol{s})$  as the probability of a subset permutation  $\bar{\pi}$  given that there are  $n_k$  elements in  $S_k$  and thus  $\bar{\pi} \in \{0,1\}^{n_k \times n}$ . Additionally, we condition on the subsets  $\boldsymbol{y}_{<k}$  and  $n_k$ , the size of subset  $S_k$ . In contrast to the distribution over permutations matrices  $p(\pi; \boldsymbol{s})$  in Equation (3), we take the product over  $n_k$  terms and have a different normalization constant  $Z_k$ . Although we induce an ordering over all elements  $i$  in Definition B.2, the probability  $p_{\boldsymbol{y}_k}$  is invariant to intra-subset orderings of elements  $i \in S_k$ .
Lemma B.6 (Probability distribution  $p_{\mathbf{y}_k}$ ). The probability distribution over subset assignments  $p_{\mathbf{y}_k}$  is given by
$$
p _ {\boldsymbol {y} _ {k}} := p (\boldsymbol {y} _ {k} \mid \boldsymbol {y} _ {<   k}; \boldsymbol {\omega}, \boldsymbol {s}) = p (n _ {k} \mid n _ {<   k}; \boldsymbol {\omega}) \sum_ {\bar {\pi} \in \Pi_ {\boldsymbol {y} _ {k}}} p (\bar {\pi} \mid n _ {k}, \boldsymbol {y} _ {<   k}; \boldsymbol {s})
$$
where  $\Pi_{\pmb{y}_k} = \{\bar{\pi}\in \{0,1\}^{n_k\times n}:\pmb {y}_k = \sum_{i = 1}^{n_k}\bar{\pi}_i\}$  and  $p(\bar{\pi}\mid n_k,\pmb {y}_{< k};\pmb {s})$  as in Lemma B.5.
Proof. We can proof the statement of Lemma B.6 as follows:
$$
\begin{array}{l} p _ {\boldsymbol {y} _ {k}} = p (\boldsymbol {y} _ {k} \mid \boldsymbol {y} _ {<   k}; \boldsymbol {\omega}, \boldsymbol {s}) \\ = \sum_ {n _ {k} ^ {\prime}} p \left(\boldsymbol {y} _ {k}, n _ {k} ^ {\prime} \mid \boldsymbol {y} _ {<   k}; \boldsymbol {\omega}, \boldsymbol {s}\right) (29) \\ = \sum_ {n _ {k} ^ {\prime}} p \left(n _ {k} ^ {\prime} \mid \boldsymbol {y} _ {<   k}; \boldsymbol {\omega}, \boldsymbol {s}\right) p \left(\boldsymbol {y} _ {k} \mid n _ {k} ^ {\prime}, \boldsymbol {y} _ {<   k}; \boldsymbol {\omega}, \boldsymbol {s}\right) (30) \\ = \sum_ {n _ {k} ^ {\prime}} p \left(n _ {k} ^ {\prime} \mid n _ {<   k}; \boldsymbol {\omega}, \boldsymbol {s}\right) p \left(\boldsymbol {y} _ {k} \mid n _ {k} ^ {\prime}, \boldsymbol {y} _ {<   k}; \boldsymbol {s}\right) (31) \\ = p \left(n _ {k} \mid n _ {<   k}; \boldsymbol {\omega}, \boldsymbol {s}\right) p \left(\boldsymbol {y} _ {k} \mid n _ {k}, \boldsymbol {y} _ {<   k}; \boldsymbol {s}\right) (32) \\ = p \left(n _ {k} \mid n _ {<   k}; \boldsymbol {\omega}\right) \sum_ {\bar {\pi} \in \Pi_ {\boldsymbol {y} _ {k}}} p (\bar {\pi} \mid n _ {k}, \boldsymbol {y} _ {<   k}; \boldsymbol {s}) (33) \\ \end{array}
$$
Equation (29) holds by marginalization, where  $n_k'$  denotes the random variable that stands for the size of subset  $S_k$ . By Bayes' rule, we can then derive Equation (30). The next derivations stem from the fact that we can compute  $n_{<k}$  if  $y_{<k}$  is given, as the assignments  $y_{<k}$  hold information on the size of subsets  $S_{<k}$ . More explicitly,  $n_i = \sum_{j=1}^{n} y_{ij}$ . Further,  $y_k$  is independent of  $\omega$  if the size  $n_k'$  of subset  $S_k$  is given, leading to Equation (31). We further observe that  $p(y_k \mid n'_k, y_{<k}; s)$  is only non-zero, if  $n'_k = \sum_{i=1}^{n} y_{ki} = n_k$ . Dropping all zero terms from the sum in Equation (31) thus results in Equation (32). Finally, by Definition B.2, we know that  $y_k = \sum_{i=\nu_k+1}^{\nu_k+n_k} \pi_i$ , where  $\nu_k = \sum_{\iota=1}^{k-1} n_\iota$  and  $\pi \in \{0,1\}^{n \times n}$  a permutation matrix. Hence, in order to get  $y_k$  given  $y_{<k}$ , we need to marginalize over all permutations of the elements of  $y_k$  given that the elements in  $y_{<k}$  are already ordered, which corresponds exactly to marginalizing over all subset permutation matrices  $\bar{\pi}$  such that  $y_k = \sum_{i=1}^{n_k} \bar{\pi}_i$ , resulting in Equation (33).
In Lemma B.6, we describe the set of all subset permutations  $\bar{\pi}$  of elements  $i\in S_k$  by  $\Pi_{\mathbf{y}_k}$ . Put differently, we make  $p(\mathbf{y}_k\mid \mathbf{y}_{< k};\boldsymbol {\omega},\mathbf{s})$  invariant to the ordering of elements  $i\in S_k$  by marginalizing over the probabilities of subset permutations  $p_{\bar{\pi}_k}$  (Xie and Ermon, 2019).
Using Lemmas B.5 and B.6, we propose the two-stage random partition  $p(Y; \omega, s)$ . Since  $Y = [\pmb{y}_1, \dots, \pmb{y}_K]^T$ , we calculate  $p(Y; \omega, s)$ , the PMF of the two-stage RPM, sequentially using Lemmas B.5 and B.6, where we leverage the PL distribution for permutation matrices  $p(\pi; s)$  to describe the probability distribution over subsets  $p(\pmb{y}_k \mid \pmb{y}_{<k}; \omega, s)$ .
Proposition 4.1 (Two-Stage Random Partition Model). Given a probability distribution over subset sizes  $p(\pmb{n};\pmb{\omega})$  with  $\pmb{n} \in \mathbb{N}_0^K$  and distribution parameters  $\pmb{\omega} \in \mathbb{R}_{+}^{K}$  and a PL probability distribution over random orderings  $p(\pi; s)$  with  $\pi \in \{0,1\}^{n\times n}$  and distribution parameters  $s \in \mathbb{R}_{+}^{n}$ , the probability mass function  $p(Y; \pmb{\omega}, s)$  of the two-stage RPM is given by
$$
p (Y; \boldsymbol {\omega}, s) = p (\boldsymbol {y} _ {1}, \dots , \boldsymbol {y} _ {K}; \boldsymbol {\omega}, s) = p (\boldsymbol {n}; \boldsymbol {\omega}) \sum_ {\pi \in \Pi_ {Y}} p (\pi ; s) \tag {34}
$$
where  $\Pi_Y = \{\pi : \pmb{y}_k = \sum_{i=\nu_k+1}^{\nu_k+n_k} \pmb{\pi}_i, k=1,\dots,K\}$ , and  $\pmb{y}_k$  and  $\nu_k$  as in Definition B.2.
Proof. Using Lemmas B.5 and B.6, we write
$$
\begin{array}{l} p (Y) = p \left(\boldsymbol {y} _ {1}, \dots , \boldsymbol {y} _ {K}; \boldsymbol {\omega}, \boldsymbol {s}\right) = p \left(\boldsymbol {y} _ {1}; \boldsymbol {\omega}, \boldsymbol {s}\right) \dots p \left(\boldsymbol {y} _ {K} \mid \left\{\boldsymbol {y} _ {j} \right\} _ {j <   K}; \boldsymbol {\omega}, \boldsymbol {s}\right) \\ = \left(p \left(n _ {1}; \boldsymbol {\omega}\right) \sum_ {\bar {\pi} _ {1} \in \Pi_ {\mathcal {Y} _ {1}}} p \left(\bar {\pi} _ {1} \mid n _ {1}; \boldsymbol {s}\right)\right) \\ \cdots \left(p \left(n _ {K} \mid \left\{n _ {j} \right\} _ {j <   K}; \omega\right) \sum_ {\bar {\pi} _ {K} \in \Pi_ {\mathbf {y} _ {K}}} p \left(\bar {\pi} _ {K} \mid \left\{n _ {j} \right\} _ {j \leq K}; \boldsymbol {s}\right)\right) (35) \\ = p \left(n _ {1}; \omega\right) \dots p \left(n _ {K} \mid \left\{n _ {K} \right\} _ {j <   K}; \omega\right) \\ \left. \cdot \left(\sum_ {\bar {\pi} _ {1} \in \Pi_ {y _ {1}}} p \left(\bar {\pi} _ {1} \mid n _ {1}; \boldsymbol {s}\right) \dots \sum_ {\pi_ {K} \in \Pi_ {y _ {K}}} p \left(\bar {\pi} _ {K} \mid \{n _ {j} \} _ {j \leq K}; \boldsymbol {s}\right)\right) \right. (36) \\ = p (\boldsymbol {n}; \omega) \left(\sum_ {\bar {\pi} _ {1} \in \Pi_ {\boldsymbol {y} _ {1}}} \dots \sum_ {\pi_ {K} \in \Pi_ {\boldsymbol {y} _ {K}}} p \left(\bar {\pi} _ {1} \mid n _ {1}; \boldsymbol {s}\right) \dots p \left(\bar {\pi} _ {K} \mid \{n _ {j} \} _ {j \leq K}; \boldsymbol {s}\right)\right) (37) \\ = p (\boldsymbol {n}; \omega) \sum_ {\pi \in \Pi_ {Y}} p (\pi \mid \boldsymbol {n}; s) (38) \\ = p (\boldsymbol {n}; \boldsymbol {\omega}) \sum_ {\pi \in \Pi_ {Y}} p (\pi ; \boldsymbol {s}) (39) \\ \end{array}
$$
B.3 Approximating the Probability Mass Function
Lemma 4.2.  $p(Y;\omega ,s)$  can be upper and lower bounded as follows
$$
\forall \pi \in \Pi_ {Y}: p (\boldsymbol {n}; \boldsymbol {\omega}) p (\pi ; \boldsymbol {s}) \leq p (Y; \boldsymbol {\omega}, \boldsymbol {s}) \leq | \Pi_ {Y} | p (\boldsymbol {n}; \boldsymbol {\omega}) \max  _ {\tilde {\pi}} p (\tilde {\pi}; \boldsymbol {s}) \tag {40}
$$
Proof. Since  $p(\pi; s)$  is a probability we know that  $\forall \pi \in \{0, 1\}^{n \times n} \quad p(\pi; s) \geq 0$ . Thus, it follows directly that:
$$
\forall \pi \in \Pi_ {Y}: p (Y; \boldsymbol {\omega}, \boldsymbol {s}) = p (\boldsymbol {n}; \boldsymbol {\omega}) \sum_ {\pi^ {\prime} \in \Pi_ {Y}} p \left(\pi^ {\prime}; \boldsymbol {s}\right) \geq p (\boldsymbol {n}; \boldsymbol {\omega}) p (\pi ; \boldsymbol {s}),
$$
proving the lower bound of Lemma 4.2.
On the other hand, can prove the upper bound in Lemma 4.2 by:
$$
\begin{array}{l} p (Y; \boldsymbol {\omega}, \boldsymbol {s}) = p (\boldsymbol {n}; \boldsymbol {\omega}) \sum_ {\pi^ {\prime} \in \Pi_ {Y}} p \left(\pi^ {\prime}; \boldsymbol {s}\right) \\ \leq p (\boldsymbol {n}; \boldsymbol {\omega}) \sum_ {\pi^ {\prime} \in \Pi_ {Y}} \max  _ {\pi \in \Pi_ {Y}} p (\pi ; \boldsymbol {s}) \\ = p (\boldsymbol {n}; \boldsymbol {\omega}) \max  _ {\pi \in \Pi_ {Y}} p (\pi ; \boldsymbol {s}) \sum_ {\pi^ {\prime} \in \Pi_ {Y}} 1 \\ = \left| \Pi_ {Y} \right| \cdot p (\boldsymbol {n}; \boldsymbol {\omega}) \max  _ {\pi \in \Pi_ {Y}} p (\pi ; \boldsymbol {s}) \\ \leq \left| \Pi_ {Y} \right| \cdot p (\boldsymbol {n}; \boldsymbol {\omega}) \max  _ {\pi} p (\pi ; \boldsymbol {s}) \\ \end{array}
$$
We can compute the maximum probability  $\max_{\pi} p(\pi; s)$  with the probability of the permutation matrix  $f_{\pi}(s)$ , which sorts the unperturbed scores in decreasing order.
B.4 The Differentiable Random Partition Model
We propose the DRPM  $p(Y; \omega, s)$ , a differentiable and reparameterizable two-stage RPM.
Lemma 4.3 (DRPM). A two-stage RPM is differentiable and reparameterizable if the distribution over subset sizes  $p(\pmb{n};\pmb{\omega})$  and the distribution over orderings  $p(\pi; s)$  are differentiable and reparameterizable.
Proof. To prove that our two-stage RPM is differentiable we need to prove that we can compute gradients for the bounds in Lemma 4.2 and to provide a reparameterization scheme for the two-stage approach in Definition B.2.
Gradients for the bounds: Since we assume that  $p(n; \omega)$  and  $p(\pi; s)$  are differentiable and reparameterizable, we only need to show that we can compute  $|\Pi_Y|$  and  $\max_{\tilde{\pi}} p(\tilde{\pi}; s)$  in a differentiable manner to prove that the bounds in Lemma 4.2 are differentiable. By definition (see Section 4.1),
$$
\left| \Pi_ {Y} \right| = \prod_ {k = 1} ^ {K} \left| \Pi_ {\boldsymbol {y} _ {k}} \right| = \prod_ {k = 1} ^ {K} n _ {k}!.
$$
Hence,  $|\Pi_Y|$  can be computed given a reparametrized version  $n_k$ , which is provided by the reparametrization trick for the MVHG  $p(\pmb{n};\pmb{\omega})$ . Further, from Equation (14), we immediately see that the most probable permutation is given by the order induced by sorting the original, unperturbed scores  $s$  from highest to lowest. This implies that  $\max_{\tilde{\pi}} p(\tilde{\pi};s) = p(\pi_s;s)$ , which we can compute due to  $p(\pi_s;s)$  being differentiable according to our assumptions.
Reparametrization of the two-stage approach: Given reparametrized versions of  $n$  and  $\pi$ , we compute a partition as follows:
$$
\boldsymbol {y} _ {k} = \sum_ {i = \nu_ {k} + 1} ^ {\nu_ {k} + n _ {k}} \boldsymbol {\pi} _ {i}, \quad \text {w h e r e} \quad \nu_ {k} = \sum_ {\iota = 1} ^ {k - 1} n _ {\iota} \tag {41}
$$
The challenge here is that we need to be able to backpropagate through  $n_k$ , which appears as an index in the sum. Let  $\alpha_k = \{0,1\}^n$ , such that
$$
(\boldsymbol {\alpha} _ {k}) _ {i} = \left\{ \begin{array}{l l} 1 & \text {i f} \nu_ {k} <   i \leq \nu_ {k + 1} \\ 0 & \text {o t h e r w i s e} \end{array} \right.
$$
Given such  $\alpha_{k}$ , we can rewrite Equation (41) with
$$
\boldsymbol {y} _ {k} = \sum_ {i = 1} ^ {n} \left(\boldsymbol {\alpha} _ {k}\right) _ {i} \boldsymbol {\pi} _ {i}. \tag {42}
$$
While this solves the problem of propagating through sum indices, it is not clear how to compute  $\alpha_{k}$  in a differentiable manner. Similar to other works on continuous relaxations (Jang et al., 2016; Maddison et al., 2017), we can compute a relaxation of  $\alpha_{k}$  by introducing a temperature  $\tau$ . Let us introduce auxiliary function  $f:\mathbb{N}\to [0,1]^n$ , that maps an integer  $x$  to a vector with entries
$$
f _ {i} (x; \tau) = \sigma \left(\frac {x - i + \epsilon}{\tau}\right),
$$
such that  $f_{i}(x;\tau) \approx 0$  if  $\frac{x - i}{\tau} < 0$  and  $f_{i}(x;\tau) \approx 1$  if  $\frac{x - i}{\tau} \geq 0$ . Note that  $\sigma(\cdot)$  is the standard sigmoid function and  $\epsilon < \epsilon / 1$  is a small positive constant to break the tie at  $\sigma(0)$ . We then compute an approximation of  $\alpha_{k}$  with
$$
\tilde {\alpha} _ {k} (\tau) = f \left(\nu_ {k}; \tau\right) - f \left(\nu_ {k - 1}; \tau\right),
$$
$\tilde{\alpha}_k(\tau) \in [0,1]^n$ . Then, for  $\tau \to 0$  we have  $\tilde{\alpha}_k(\tau) \to \alpha_k$ . In practice, we cannot set  $\tau = 0$  since this would amount to a division by 0. Instead, we can apply the straight-through estimator (Bengio et al., 2013) to the auxiliary function  $f(x;\tau)$  in order to get  $\tilde{\alpha}_k \in \{0,1\}^n$  and use it to compute Equation (42).