Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
88,552 characters · 17 sections · 93 citation commands
Decision Theory for the Archetype Discovery Problem
\onehalfspacing
\doublespacing
A researcher has access to a function $\phi$ that describes the variation of $N$ heterogeneous effects of a policy of interest over a discrete set of observable characteristics, $x \in \mathcal{X}$. The researcher would like to communicate this information to a policymaker that wishes to assess the impact of the policy over a group of heterogeneous individuals. Unfortunately, when $N$ is large, communicating $\phi$ to the policymaker could be challenging. The policymaker might prefer a simpler summary of these effects. For example, it is possible that the policymaker finds it simpler to interpret the Sorted Group Average Treatment Effects (GATES) of Victor25, where the policy effects in $\phi$ are first sorted to obtain a few quantiles and then $\phi$ is summarized by reporting averages of policy effects over the groups defined by such quantiles. The researcher thus faces a critical trade-off when communicating the policy effects: the reports should contain as much information as in the original function $\phi$ (that is, the reports should exhibit high fidelity), but should also be easy to parse for the policymaker (the reports should be less complex than the original function $\phi$).
In recent work, breza2025generalizability develop a useful framework to think about how to navigate the aforementioned trade-off between high fidelity and low complexity. In particular, they consider the problem of how to partition the researcher's heterogeneous policy effects into $K$ groups of “individual and environmental” characteristics within which treatment effects are “predictively stable.” breza2025generalizability refer to the groups of characteristics as archetype sets, and to the problem of finding these groups as the problem of archetype discovery. In an important departure from existing literature, they allow the researcher to pay a cost to admit ignorance for some groups.
In this paper, we apply statistical decision theory to the archetype discovery problem. We focus on the case in which the researcher is not allowed to admit ignorance and must summarize the policy effects for every observable characteristic. Since there are different recommendations in the literature regarding how to present and summarize heterogeneous policy effects, decision theory can guide their evaluation.\footnote{In addition to the GATES of Victor25 and the generalizability aware predictions of breza2025generalizability, other suggestions in the literature include the endogenous stratification strategy studied by abadie2018endogenous, the strategies based on decision trees in AtheyImbens2016,WagerAthey2018, and the rashomon partitions of venkateswaran2024robustly.} Using statistical decision theory necessitates the specification of at least three ingredients: the menu of actions available to the researcher, the consequences of these actions as a function of a potentially unknown state of the world (i.e., the loss function), and a statistical model which captures how the data distribution changes at each possible unknown state of the world. In terms of the action space, we focus on the case in which the researcher is only allowed to communicate a function $\bar{\phi}$ that takes at most $K$ values. This choice of action space ensures that the function reported by the researcher has low complexity relative to $\phi$, since $K$ will typically be small and, in particular, considerably smaller than the number of values that $\phi$ could take. Throughout the paper, we assume that the researcher takes $K$ as exogenously given. For the loss function, we assume that the main criterion used by the researcher in order to guide the construction of the report $\bar{\phi}$ is a typical weighted mean-squared error criterion. In particular, the loss associated to reporting the function $\bar{\phi}$ to summarize the original function $\phi$ takes the form:
where $p(x)$ denotes a probability mass function over the discrete set $\mathcal{X}$ that is assumed to be known and chosen by the researcher. The loss function in (ref) encourages reports with high fidelity by penalizing the differences between $\phi$ and $\bar{\phi}$.
Our first result concerns the oracle choice of $\bar{\phi}$: that is, we are interested in characterizing the researcher's report that minimizes the loss function in (ref), assuming that the policy effects in $\phi$ are known. Our Theorem (ref) shows that the oracle choice of $\bar{\phi}$ can be constructed by a procedure analogous to GATES. In the oracle procedure characterized by our theorem, the $N$ policy effects in $\phi$ are first sorted in increasing order and $K$ groups are then constructed by solving a weighted version of the popular $K$-means problem of mcqueen1967some and lloyd1982least. An important observation here is that the $K$ groups are based on clustering the $N$ scalar policy effects, not on the covariates, as in the recent work of kim2026causal. It then follows that the main conceptual difference between the oracle solution and the popular GATES procedure is that in the oracle solution the groups are obtained by weighted $K$-means clustering of the policy effects, instead of relying on $K$ equally-spaced quantiles. Moreover, since the $K$ groups are defined by clustering the policy effects, the $k$-th archetype set in the oracle problem simply consists of all covariate values that map to a policy effect that belongs to the $k$-th cluster. The oracle report $\bar{\phi}$ is constant over each cluster, and the reported value is the cluster-conditional mean of $\phi$.
Before presenting the other results in the paper, it is worthwhile making two additional comments about the oracle choice of $\bar{\phi}$. First, a common critique of the GATES procedure (and of any other endogenous stratification procedure that creates groups based on outcomes) is that the resulting archetype sets could be “uninterpretable in terms of the original covariates, making it necessary to post-process the treatment effect distribution to gain insights about particular covariates”; see the discussion in Section H.4 of venkateswaran2024robustly. The proof of our Theorem (ref) shows that the interpretability (or lack thereof) of the oracle solution $\bar{\phi}$ crucially depends on the original policy effects $\phi$. It is perhaps convenient to explain this in the context of a simple example. Suppose that the original policy effects are described by a decision tree and that the covariate space is a finite subset of the unit square $[0,1]^2$. The different values of the policy effects partition the unit square in terms of axis-aligned rectangles. Interestingly, our results show that the archetype sets associated to the oracle solution will correspond to a coarsening of this partition. This means that $\bar{\phi}$ inherits some of the interpretability in $\phi$ (unions of axis-aligned rectangles), without having to impose interpretability as an additional constraint in the oracle problem.
Second, since the $K$ groups in the oracle solution are obtained by solving a weighted $K$-means clustering problem of $N$ scalar policy effects (as opposed to solving a clustering problem involving a potentially high-dimensional vector of covariates), the exact solution to the oracle problem can be obtained in a computationally convenient way. It is known---see bruce1965optimum, wu1989kn,wu1991optimal, wang2011ckmeans, gronlund2017fast---that the weighted $K$-means clustering problem of $N$ sorted scalars can be solved in $O(KN^2)$ time using $O(KN)$ space by means of a simple dynamic programming algorithm that---for the sake of exposition---we present in Section (ref). It is important to note that the key feature that enables the use of the dynamic programming algorithm is that the $K$ clusters of the policy effects are contiguous: whenever two policy effects $\phi_n<\phi_{n'}$ are assigned to the same cluster, then any other policy effect that is ordered in between them will also be assigned to the same cluster. This means that the archetype sets in the oracle solution can be thought of as arising from a “squashing” of the level curves of the original function $\phi$.
While our analysis of the oracle solution to the archetype discovery problem is helpful to understand how to summarize policy effects if one must, in practice, the function $\phi$ will need to be estimated from the data. Thus, we use statistical decision theory to understand if the variant of the GATES procedure that we have discussed so far is still a reasonable way of using sample data for archetype discovery. To this end, we first consider a Bayesian statistical decision problem, where we assume that the researcher observes a dataset $z \in \mathcal{Z}$ that is informative about $\phi$. The researcher has a (potentially nonparametric) statistical model $\{P_{\theta}\}_{\theta \in \Theta}$, where the index $\theta$ includes $\phi$, but potentially other nuisance parameters. We endow the researcher with a prior $\pi$ over $\Theta$.
Our second result, Theorem (ref), shows that---under the assumption that the model's parameter $\theta$ only enters the loss function through $\phi$---it is optimal to solve the archetype discovery problem by applying the dynamic programming algorithm to the posterior mean of $\phi$. This means that if the policy effects in $\phi$ are Conditional Average Treatment effects estimated from experimental or observational data, one could use the Bayesian regression tree models for causal inference of hahn2020bayesian based on the seminal work of hill2011bayesian. The archetype sets can then be formed by clustering the posterior mean estimates of the conditional treatment effects. The resulting procedure is still analogous to GATES, but the key difference now is that the proxies for the treatment effects are based on a Bayesian posterior mean estimator, as opposed to a machine learning proxy.
Our third result, Theorem (ref), dispenses the use of the prior $\pi$, and focuses on a researcher that is interested in minimizing worst-case risk. In order to analyze this problem, we need to be more explicit about the statistical model. To this end, we assume that the researcher has access to an estimator of the form
where we allow for distributions $P$ for which the marginals of $u_x$ have mean zero, variance one, and are subgaussian (with an optimal variance proxy of at most one). We treat $\sigma_x$ and $I$ as known. The hyperparameter $I$ plays a role analogous to the sample size: a large value of $I$ is interpreted as having a more precise estimator of $\phi(\cdot)$. We also note that we allow the error terms, $u_x$, to be correlated across the different values of the covariates $x \in \mathcal{X}$. Theorem (ref) shows that applying the dynamic programming algorithm to the sorted values of $\hat{\phi}$ is approximately minimax (in a sense we make precise) provided $I$ is large enough. We also show in Theorem (ref) that the minimax regret converges to zero at a rate that is at most of order $\sqrt{\log(|\mathcal X|)/I}$.
Finally, we consider the case in which---in addition to communicating the function $\bar{\phi}$ to the policymaker---the researcher can also provide a set of covariate values at which it is better to admit ignorance or abstain. We adopt the same loss function as in breza2025generalizability, and just like in their framework, we consider a hyperparameter that measures how costly it is for the researcher to admit ignorance. We make a restriction on the type of abstention that the researcher can recommend to the policy maker. In particular, we require the abstention to respect archetype sets: if the researcher recommends abstention for one covariate value in the $k$-th archetype set, then the researcher must necessarily recommend abstention for all the covariate values in such set. We think this is a reasonable restriction (consistent with the idea that the policymaker has some limited ability to parse complex functions). We show that this restriction also leads to reports $\bar{\phi}$ that cluster units based on their policy effects. If we further require the clusters to be contiguous (in a sense we make precise), we show there is a simple extension of our dynamic programming algorithm (that modifies the flow cost of the dynamic programming problem) that solves the archetype discovery problem with abstention.
The rest of the paper is organized as follows. Section (ref) introduces the formal framework and characterizes the oracle solution to the archetype discovery problem (the best way of summarizing the information in $\phi$ when this function is known). Section (ref) considers the case in which the policy effects of interest are unknown and studies data-driven rules under the average-risk and worst-case-risk criteria. Section (ref) presents a simple numerical example to illustrate our results. Section (ref) discusses different extensions of the main results: in particular, we consider the archetype discovery problem with abstention, and also discusses alternative algorithms for solving the weighted $K$-means clustering problem. Section (ref) revisits the application discussed in Victor25: an experiment with the government of Haryana in North India designed to analyze the effects of a policy bundle that provided different incentives for immunization across several villages. We use this application to illustrate the differences between archetype sets constructed via weighted $K$-means clustering and archetype sets constructed using quantiles of the policy effects (as in the GATES procedure). Proofs of the main results are collected in Appendix (ref). Additional results and supporting material are collected in Appendix (ref).
A researcher has access to a mapping $\phi: \mathcal{X} \rightarrow \mathbb{R}$ that describes the effects of a policy of interest as a function of a discrete set of observable characteristics. The researcher would like to communicate the policy effects contained in $\phi$ to a policymaker, but the policymaker would prefer a simpler summary of these policy effects.
Let $\{\phi_1, \ldots, \phi_N\}$ denote the $N \leq | \mathcal{X}|$ different values that the function $\phi$ takes. While the policymaker has a hard time parsing the information contained in $\phi$, the policymaker is comfortable processing and interpreting a function $\bar{\phi}: \mathcal{X} \rightarrow \mathbb{R}$ that takes only $K$ values, with $K\ll N$. The researcher is thus interested in finding a function $\bar{\phi}$ that provides a policy-relevant summary of the information in $\phi$.
We refer to the set of characteristics associated to the $k$-th value of $\bar{\phi}$ as the $k$-th archetype set. The problem of finding the archetype sets (and their corresponding values) was recently introduced by breza2025generalizability. We follow their terminology and refer to this problem as the archetype discovery problem.
Let $p(x)$ denote a probability mass function over the discrete set $\mathcal{X}$. We follow breza2025generalizability and assume that the main criterion used by the researcher in order to guide the construction of the summary $\bar{\phi}$ is a typical weighted squared loss. In particular, the loss of using $\bar{\phi}$ to summarize the original function $\phi$ takes the form:
Throughout, we assume, without loss of generality, that $p(x)>0$ for all $x \in \mathcal{X}$. We also assume that the researcher takes $K$---the number of values that the policymaker can comfortably parse---as given and tries to minimize (ref) over all functions that take only $K$ values. More formally, let $\bar{\Phi}_{K}$ denote the set of all functions that map the set $\mathcal{X}$ to the real line and that take at most $K$ values; that is \[ \bar{\Phi}_{K} \equiv \left\{ \bar{\phi}:\mathcal{X} \rightarrow \mathbb{R} \: \mid \: | \bar{\phi}(\mathcal{X})| \leq K \right\}, \] where $\bar{\phi}(\mathcal{X})$ denotes the image of the set $\mathcal{X}$ under the function $\bar{\phi}$. We say that $\bar{\phi}^*$ is an oracle solution to the archetype discovery problem if
Let $\bar{\phi}^*_{k}$ denote the $k$-th value of the oracle solution $\bar{\phi}^*$. Define the oracle $k$-th archetype set associated to the oracle solution $\bar{\phi}^*$ as
Our first result shows that the oracle solution to the archetype discovery problem can be obtained by optimally grouping the $N$ values $\{ \phi_1, \ldots, \phi_N\}$ into $K$ clusters. We will show that this can be done by solving the one-dimensional version of the classical $K$-means problem in mcqueen1967some and lloyd1982least. The signal processing literature refers to this problem as optimum $K:N$ quantization; see wu1989kn,wu1991optimal and the references therein.
In order to establish our first result, assume (without loss of generality) that the values of $\phi$ have been sorted and that $\phi_1< \ldots < \phi_{N}$. Define a $K$-clustering of the values $\phi_1 < \ldots <\phi_{N}$ as a surjective function $c:\{1,\ldots, N\} \rightarrow \{1,\ldots, K\}$. We interpret $c(i)$ as the cluster assigned to the value $\phi_i$. The $k$-th cluster associated to the clustering function $c(\cdot)$ can then be denoted as $C_k \equiv \{\phi_i \in \{\phi_1, \ldots, \phi_N\} \: \mid \: c(i) = k \}$.
Given a clustering function, let $\mu_{k}(c)$ denote the (conditional) mean of $\phi$ within cluster $k$; that is
Consider then the problem of clustering the values $\{\phi_1, \ldots, \phi_{N}\}$ into $K$ clusters in order to minimize weighted squared loss:
where the minimum is taken over all $K$-clustering functions.
Theorem (ref) shows that the archetype discovery problem in (ref) can be solved in the following way. The researcher first needs to sort the $N$ different values of $\phi$. Then, the researcher optimally clusters these values into $K$ groups---by solving (ref). Note that the characteristics $x$ do not enter this process: the clustering is done over $\{\phi_1, \ldots, \phi_N\}$ not over $\mathcal{X}$. The solution to the archetype discovery problem, however, is a $K$-valued function $\bar{\phi}^*$ defined over $\mathcal{X}$ and the theorem shows that the $\bar{\phi}^*$ can be constructed as follows. For a given $x \in \mathcal{X}$ we first retrieve the index $i \in \{1, \ldots, N\}$ associated with the value that the original function $\phi(x)$ is assigned to. For example, suppose that when evaluated at $\tilde{x}$, we have $\phi(\tilde{x})= \phi_3$. Then $i(\tilde{x})=3$. Once we know that $\phi(x)$ takes the value $\phi_{i(x)}$ we look for the cluster to which $i(x)$ was assigned. Under the optimal clustering, this value is $c^*(i(x))$. Our result says that the solution to the archetype discovery problem can be obtained directly from the $K$ optimal clusters of the set $\{\phi_1, \ldots, \phi_N\}$ by simply reporting the (conditional) cluster mean associated to $c^*(i(x))$. That is $\bar{\phi}^*(x) = \mu_{c^*(i(x))}(c^*)$.
We note that the solution to the archetype discovery problem described in Theorem (ref) is analogous to the Sorted Group Average Treatment Effects (GATES) procedure described in Victor25, assuming away the (machine learning) estimation of the heterogeneous effects. As explained by Victor25 (see Comment 3.6, p. 1135), the groups created by GATES are “based upon actual predicted treatment effect". This means that the “heterogeneity groups” created by GATES are groups induced by the ML proxy predictor (see p. 1128 in their paper). When there is no sampling uncertainty, this is analogous to create groups based on the values of the policy effects contained in $\phi$.
The key difference between the oracle solution in Theorem (ref) and the GATES procedure, is that the policy effects in $\phi$ are obtained by weighted $K$-means clustering of the $N$ heterogeneous policy effects, instead of relying on $K$ equally-spaced quantiles.\footnote{For instance, in their main application Victor25 estimate the GATES by quintiles of the ML proxies. See their Figure 4 in Section 6.2.}
Finally, it is worthwhile to give a brief interpretation of our Theorem (ref) in terms of the well-known signal processing literature on optimum quantization; see wu1989kn,wu1991optimal. Using the terminology from this literature, our Theorem (ref) says that the optimal low-complexity report of $\phi$ is the optimum $K$-point scalar quantizer of the push-forward distribution of $X \sim p(\cdot)$ by $\phi(\cdot)$, with archetype sets given by inverse images of quantizer cells.
It is known that an exact solution to the one-dimensional clustering problem in (ref) can be found using dynamic programming. The algorithm we present below is based on the seminal work on optimum quantization of bruce1965optimum. We provide matlab and python scripts that implement such algorithm.\footnote{wang2011ckmeans provide an implementation in R.} bruce1965optimum's algorithm performs well in our applications (and runs in a fraction of a second), but we note that there are other algorithms for finding an exact solution to the one-dimensional clustering problem in (ref) that can improve both the runtime and the space requirements of the dynamic programming suggested in this paper. See gronlund2017fast and the references therein. We decided to focus on bruce1965optimum's algorithm because of its simplicity.
Dynamic Programming algorithm. We are interested in grouping $\phi_1 < \ldots <\phi_{N}$ into $K$ clusters using the weighted $K$-means objective function in (ref). The first key observation is that we can focus on clustering functions $c:\{1,\ldots, N\} \rightarrow \{1,\ldots, K\}$ that produce contiguous clusters: if $\phi_n < \phi_{n'}< \phi_{n{''}}$ and $c(n)=c(n'')=k$, then $c(n')=k$. Note that any such clustering function can be characterized by a collection of $K-1$ ordered integers $q_0 \equiv 0 < q_1 < \ldots < q_{K-1}<q_{K} \equiv N$ where $c(i)=k$ if and only if $q_{k-1} < i \leq q_{k}$. The $k$-th cluster is thus given by $\{ \phi_i \: \mid \: q_{k-1} < i \leq q_{k} \}$. Setting $q_0 \equiv 0$ and $q_K \equiv N$ guarantees that the $K$ clusters are well defined.
Note that cardinality of the set of all contiguous clusters is given by the combination formula of $(N-1)$ choose $(K-1)$. A naive algorithm that works via enumeration becomes intractable for modest case and $K \ll N$ and $N$ is large. For example, if $N=100$ and $K=6$, the total number of elements is 715,231,440.
A key observation that allows for a dynamic programming algorithm to find the optimal clustering scheme is the following. Let $\widehat{q}^K_{N} = (q_1, \ldots, q_{K-1})$ be the vector that contains the $K-1$ ordered thresholds that define the optimal $K$ clusters. Lemma 1 in wu1989kn shows that in the case of multiplicity of solutions, we can always select $\widehat{q}^K_{N}$ to be the solution with the largest thresholds (the set of solutions is totally ordered under the usual vector ordering). Let $q_{t}$ denote the $t$-th ($\leq K-1$) entry of the vector $\widehat{q}^K_{N}$. As discussed above the threshold $q_t$ defines the $t$-th cluster for $\phi_1< \ldots< \phi_N$. Consider the points $\phi_1< \ldots< \phi_{q_t}$. Lemma 2 in wu1989kn shows that $(q_1, q_2, \ldots, q_{t-1})$ still characterize the optimal way of clustering the $q_t$ points $\phi_1< \ldots< \phi_{q_t})$ into $t$ clusters.
This means that the clustering problem admits a recursive representation. For any $1 \leq k \leq K$ and $k \leq n \leq N$, let $D(k,n)$ denote the weighted mean-squared error of optimally grouping $\phi_1< \ldots< \phi_n)$ into $k$ clusters. For $n \leq n' \leq N$, let $C(n,n')$ be the squared-error of assigning points $\phi_{n} \leq \ldots \leq \phi_{n'}$ into one cluster. That is \[ C(n,n') \equiv \sum_{i=n}^{n'} p_i \left( \phi_i-\mu_{n,n'} \right)^2, \quad \textrm{ where } \mu_{n,n'} \equiv \frac{\sum_{i=n}^{n'} p_i \phi_i}{\sum_{i=n}^{n'} p_i}. \] Note that Lemma 2 in wu1989kn implies \[ D(k,n) = \min_{a \in \{k,\ldots, n\}} C(a,n) + D(k-1,a-1). \] This defines a dynamic programming problem in the sense of bradley1977applied. The state of the problem is a vector $(k,n)$ where $k \leq n$, $1 \leq k \leq K$, and $k \leq n \leq N$. The state can be interpreted as making reference to the $n$ points $\phi_1< \ldots< \phi_n$ that have not been clustered and that need to be grouped into $k$ clusters. The action (or decision) in each state, denoted by $a \in \{k,\ldots, n\}$, refers to the points $\{\phi_a , \ldots, \phi_n\}$ that will be assigned to the right-most cluster. The flow payoff (or more precisely, the flow cost) is simply $C(a,n)$. The transition function maps a state $(k,n)$ and an action $a$ into the new state $(k-1,a-1)$, where we need to cluster $\{\phi_1, \ldots, \phi_{a-1}\}$ into $k-1$ clusters.
It is known that there are other slightly more complicated algorithms that can improve the time and space requirements of the baseline dynamic programming algorithm; see Section 1.2 in gronlund2017fast. In the numerical exercises where $N$ is in the thousands and $K$ is less than or equal to $10$, the runtime of the dynamic programming algorithm we have just discussed is almost a third of a second in a personal laptop, and therefore we recommend this procedure.
It is also useful to contrast this exact one-dimensional solution of the $k$-means clustering problem with approximate algorithms commonly used for multidimensional $k$-means problems. We note that in our framework there is no need to use heuristic algorithms to solve the $K$-means problem. The popular iterative procedures of lloyd1982least-forgy1965cluster, and hartigan1979algorithm have been applied to other clustering problems in econometrics; e.g., bonhomme2015grouped and bonhomme2019distributional. However, as noted by gronlund2017fast, even for clustering problems in one dimension, “Lloyd’s algorithm does not necessarily compute the optimal clustering, it is only a heuristic.”
In the previous section we described an oracle (or population) version of the archetype discovery problem, where we assumed that the researcher knew the function $\phi: \mathcal{X} \rightarrow \mathbb{R}$ describing the effects of the policy of interest. In this section we consider the case in which $\phi$ is potentially unknown to the researcher, but we assume there is a dataset available that is informative about this parameter. The question of interest is how to use the data to provide a summary $\bar{\phi}(\cdot)$ to the policymaker.
Suppose that the researcher observes a dataset $z \in \mathcal{Z}$. As usual, we view the data as the realization of a $\mathcal{Z}$-valued random variable. The researcher has a statistical model $\{P_{\theta}\}_{\theta \in \Theta}$, where the parameter $\theta$ includes $\phi$ and potentially other nuisance parameters. But throughout this section, we maintain the assumption that the weights $p(\cdot)$ are known. These weights encode the populations or policy weights according to which approximation errors are evaluated.
Just as we did before, we assume that the researcher takes $K$---the number of values that the policymaker can comfortably parse---as given. The main goal of the researcher is to provide a data-driven summary of the function $\phi$. The loss function that the researcher uses takes the form \[ L(\bar{\phi},\theta) = L(\bar{\phi};\phi,p) \;\equiv\; \sum_{x\in\mathcal{X}} p(x)\Big(\phi(x)-\bar{\phi}(x)\Big)^2. \] This means that we are assuming that the model's parameter $\theta$ only enters the loss function through $\phi$. We denote any data-driven summary as a function $d: \mathcal{Z} \rightarrow \bar{\Phi}_{K}$. We use the standard terminology in statistical decision theory---see, for example, Ferguson67---and refer to $d$ as the researcher's decision rule. The risk of a decision rule at $\theta$ is defined as
Consider first the case where the researcher has a prior distribution $\pi$ over the parameter space $\Theta$. And consider the problem of finding a decision rule that minimizes average risk: $\mathbb{E}_{\theta \sim \pi}[R(d,\theta)]$. It is known that such a decision rule can be found by minimizing posterior loss (c.f., berger1985statistical, Result 1, p. 159); that is, for each $z \in \mathcal{Z}$, set $d(z)$ to be any summary function $\bar{\phi}$ that solves the minimization problem
Our next result shows that the posterior loss minimization problem in (ref) can be solved by applying the dynamic programming algorithm to the posterior mean of $\phi$.
This simple result shows that in order to find the decision rule $d^*$ that minimizes average risk it suffices to solve an oracle archetype discovery problem where we pretend that the posterior mean function, $\mathbb{E}_{\phi \sim \pi}[\phi(x) \: | \: z]$, is the true function $\phi$. The results in the previous section then imply that if we define \[ \widehat{\phi}(x) \equiv \mathbb{E}_{\phi \sim \pi}[\phi(x) \: | \: z] \] and sort its image as $\widehat{\phi}_1< \ldots < \widehat{\phi}_N$ (where $N \leq |\mathcal{X}|$), we can use the dynamic programming algorithm to optimally group $\widehat{\phi}_1< \ldots < \widehat{\phi}_N$ in $K$ clusters.
We think this result is interesting for two reasons. First, Theorem (ref) provides a simple (and principled) approach to deal with the fact that, in applications, the function $\phi$ will typically be unknown. Second, the proof of Theorem (ref) shows that $d^*(z)$ will be measurable with respect to the $\sigma$-algebra generated by $\widehat{\phi}$. This means that if the posterior mean function has a special structure in its domain (for example, if $\widehat{\phi}$ is a decision tree over $\mathcal{X}$), then archetype sets corresponding to $d^*(z)$ will be a coarsening over the partitions in $\mathcal{X}$ induced by the function $\widehat{\phi}$.
Consider now the problem of finding the decision rule that minimizes worst-case risk. That is, we are interested in solving the problem \[ \adjustlimits \inf_{d} \sup_{\theta \in \Theta} R(d,\theta), \] where the risk function is defined in (ref), and $\theta$---which includes $\phi$---will denote the parameter of the statistical model used for the available data.
In order to make progress with the minimax problem, we focus on the case in which the available data consist of an estimator of the unknown function $\phi$. More precisely, suppose that we have an estimator of the form
We define the parameters of this statistical model to be $\theta \equiv (\phi,P)$ and we treat $\sigma_x$ and $I$ as known. The hyperparameter $I$ plays a role analogous to the sample size: a large value of $I$ is interpreted as having a more precise estimator of $\phi(\cdot)$. We also note that we are allowing for the error terms, $u_x$, to be correlated across the different values of the covariates $x \in \mathcal{X}$.
Fix an arbitrary constant $B>0$ and define the parameter space $\Theta$ to be any subset of
The set $\Theta(B)$ only contains functions $\phi$ that are bounded in absolute value by $B$. In addition, the set $\Theta(B)$ only allows for distributions $P$ for which the marginals of $u_x$ have mean zero, variance one, and are subgaussian; e.g. see definition in rigollet2015, Vershynin2018, rivasplata2012 (with an optimal variance proxy of at most one).
We are not aware of how to find the minimax rule when $\Theta = \Theta(B)$, or even when $\Theta$ is a strict subset of $\Theta(B)$ that only allows for error terms that are independent standard Gaussian random variables. Instead of insisting in finding an exact solution, we search for a solution that approximately solves the minimax problem when $I$ is large enough.
To this end, let $\bar{\Phi}_K(B) \equiv \{ \bar{\phi} \in \bar{\Phi}_K \: | \: \bar{\phi}(x) \in [-B,B]\}$. Define the plug-in decision rule, $d_{\textrm{plug-in}}$, to be any decision rule such that
Let $\mathcal{D}(B)$ be the set of all decision rules that map $\widehat{\phi}$ to $\bar{\Phi}_{K}(B)$. The minimax value of interest---as a function of $I$ and the parameter space $\Theta \subseteq \Theta(B)$---is
denote the finite-sample minimax benchmark. The following theorem shows that the plug-in rule is \(\varepsilon\)-minimax (in the sense of Ferguson67, Chapter 1, Definition 4, p.33) for all sufficiently large values of \(I\).
We finalize the decision-theoretic analysis of the archetype discovery problem by analyzing the worst-case regret criterion. We define the regret loss of an action $\bar{\phi} \in \bar{\Phi}_{K}(B)$ to be:
As usual, the regret loss above captures the excess loss of an action relative to the oracle solution of the archetype discovery problem.
Theorem (ref) shows that---even if the cardinality of $\mathcal{X}$ grows as a function of $I$---the worst-case regret of the optimal decision rule in the archetype discovery problem vanishes to zero as $I$ grows large (provided $|\mathcal{X}|$ does not grow too quickly and $\bar{\sigma}$ remains bounded). Thus, our theorem provides an upper bound on the minimax regret rate associated to the archetype discovery problem.
Theorem (ref) is established by showing that the worst-case regret of the plug-in rule is bounded by the right-hand side of (ref). That is, \[\sup_{\theta \in \Theta(B)} \mathbb{E}_{\widehat{\phi} \sim (\phi,P)} \left[ \mathcal{L}\left( d_{\textrm{plug-in}}( \widehat{\phi});\phi,p \right) \right] \leq 8B \bar\sigma\, \sqrt{\frac{2 \log(2|\mathcal{X}|)}{I}}. \] Incidentally, this means that the worst-case regret of the plug-in rule will tend to be small, provided $\mathcal{X}$ is not large relative to $I$.
A natural question to ask is whether there exists a matching lower bound for the worst-case regret. We think that it might be possible to give a positive answer to this question, at least in the case in which the dimension of $\mathcal{X}$ is fixed. In particular, we think that a lower bound for minimax regret of order $I^{-1/2}$ could be derived using a similar argument to the one used in the proof of Theorem 1 in bartlett1998, which provides a minimax lower bound for the regret of empirically designed vector quantizers.
In this section we present a simple example to illustrate how the solutions of the archetype discovery problem can be obtained by clustering the values of the function $\phi$.\footnote{Appendix (ref) presents an additional illustrative example using data from the Atlantic Causal Inference Conference (ACIC) 2016.} Our goal is twofold. First, we want to provide the reader a concrete sense of the computational cost associated to the use of the dynamic programming algorithm for clustering the $N$ different values of $\phi$. Second, we would like to suggest different ways of visualizing the archetype set and the values of the function $\bar{\phi}$. A distinction that is useful to keep in mind while reading this section is that there is usually a difference between the cardinality of the domain of $\phi$ (denoted $|\mathcal{X}|$) and the cardinality of the image of $\phi$ (that is, the number of different values that $\phi$ takes, denoted by $N$). With this distinction in mind, we present tables and figures with information on both \(|\mathcal{X}|\) and \(N\). We assume throughout that $p(x)=1/|\mathcal{X}|$ for every $x \in \mathcal{X}$. Thus, the optimized oracle loss is \[ L_{oracle}(c) \equiv \sum_{i=1}^{N} p_i \bigl(\phi_i-\mu_{c(i)}\bigr)^2 \qquad \text{where} \qquad p_i=\frac{|\{x \in \mathcal{X}\: | \: \phi(x)=\phi_i\}|}{|\mathcal{X}|}. \]
Consider the function \[ \phi(x_1,x_2)=\exp\!\left(-(x_1^2+x_2^2)\right), \qquad (x_1,x_2)\in[-1,1]^2, \] evaluated on an equally spaced \(300\times 300\) grid over \([-1,1]^2\). Hence \(|\mathcal X|=90{,}000\), \(N=7{,}401\), and \(K=10\).
We find the oracle solution to the archetype discovery problem using bruce1965optimum's dynamic programming algorithm described in Section (ref). The end-to-end computational complexity of this algorithm is known to be \[ O\!\bigl(|\mathcal X|\log |\mathcal X| + K N^2\bigr). \] The first component comes from sorting the values of $\phi(\mathcal{X})$, which requires computations of order \(O(|\mathcal X|\log |\mathcal X|)\). The second component is the cost of the dynamic-programming routine, which is known to be of order \(O(KN^2)\). The exact dynamic program has runtime of \(0.3688\) seconds, and the value of the oracle solution is \(L_{oracle}=0.000575\).
Panel a) of Figure (ref) below reports the archetype sets (each in different color) obtained by the dynamic programming algorithm. As discussed in Section (ref), the archetype sets inherit the structure of the level sets of the original function \(\phi\) (which, in this example, are concentric rings). Panel a) presents the archetype sets. Panel b) reports the sorted values of the image of $\phi$, where each different color represents a different cluster. The horizontal lines depict the within cluster means, which correspond to the values of the function $\bar{\phi}$ for each $x$ that leads to values of $\phi$ in that cluster.
Figure (ref) complements Figure (ref) by showing the within-cluster variance and the share of covariate values contained in each cluster. The color coding for each of the bars is the same as the one used to depict the different clusters in Figure (ref). Note that the clusters are indexed according to the value of their within-cluster mean (the first cluster, containing 6.8% of the observations, has the lowest within-cluster mean).
Consider now the case in which---in addition to communicating the function $\bar{\phi}$ to the policymaker---the researcher can also provide a set of covariate values at which, instead of making a prediction, it is better to admit ignorance or abstain. We follow breza2025generalizability and model abstentions as a function $\pi: \mathcal{X} \rightarrow \{0,1\}$, where $\pi(x)=0$ is interpreted as a suggestion (from the researcher to the policymaker) to avoid using the function $\bar{\phi}$ for making predictions at $x$. We adopt the same loss function as in breza2025generalizability:
where $\sigma^2$ denotes the cost of abstention.
In order to solve the oracle problem we impose two mild restrictions on the functions $\bar{\phi}$ and $\pi$. First, we assume that the researcher is only allowed to report a summary $\bar{\phi}$ that belongs to the set $\bar{\Phi}^*_K(A_K,G^{\phi})$ defined in Equation (ref).\footnote{That is, the function $\bar{\phi}$ is required to be measurable by the smallest $\sigma$-algebra that makes $\phi$ measurable.} As discussed in Remark (ref), this restriction does not entail any loss of generality when abstention is not allowed and we impose it to have a more direct comparison with Theorem (ref).
Second, we make a restriction on the type of abstention that the researcher can recommend to the policy maker. Define \[ \Pi_K(\bar\phi) \equiv \{ \pi: \mathcal{X} \to \{0,1\} \: | \: \bar\phi(x)=\bar\phi(x') \implies \pi(x)=\pi(x')\}.\] Note that an abstention function $\pi \in \Pi_K(\bar\phi)$ respects the archetype sets defined by $\bar{\phi}:$ if the researcher recommends abstention for one covariate value in the $k$-th archetype set, then the researcher must necessarily recommend abstention for all the covariate values in such set. We think this is a reasonable restriction (consistent with the idea that the policymaker has some limited ability to parse complex functions).
Define the oracle solution to the archetype discovery problem with $\bar{\phi}$-abstention as
Note that when $\sigma^2$ is large enough---for example, when $\sigma^2 \geq (\phi_N-\phi_1)^2$---a solution to the problem in (ref) can be obtained by solving the usual $K$-means clustering problem in (ref).\footnote{In fact, Popoviciu's inequality Popoviciu1935,BhatiaDavis2000 implies that a weaker sufficient condition is $\sigma^2 \geq (\phi_N-\phi_1)^2/4$.} However, it is relatively straightforward to provide examples in which solving (ref) (and declaring abstention whenever the cluster variance is larger than $\sigma^2$) is not optimal. For instance, suppose that we want to cluster the values $\{0,2,3,5\}$ into $K=2$ clusters, and suppose further that the weights are uniform. Assume also that the cost of abstention is $\sigma^2=1$. Algebra shows that the clusters $\{0,5\}$ and $\{2,3\}$ (with abstention for the elements in the first cluster) evaluate to a lower value of the objective function in (ref) than the $K$-means clustering solution in (ref), which consists of the contiguous clusters $\{0,2\}$ and $\{3,5\}$. To see this, note that the objective function in (ref) for clusters $\{0,5\}$ and $\{2,3\}$ equals
But the objective function in (ref) for clusters $\{0,2\}$ and $\{3,5\}$ equals 0. This means that declaring abstention for values $\{0,5\}$ and clustering together $\{2,3\}$, is better than using the $K$-means clusters $\{0,2\}$ and $\{3,5\}$ (both of which leave the researcher indifferent between using the reported cluster means or abstaining).
Algorithms for archetype discovery with abstention. It is not clear to us that---without further restrictions---there exists an algorithm for solving the oracle archetype discovery problem with abstention in (ref) that runs in polynomial time in $(K,N)$. When the clusters are not required to be contiguous, the problem seems to be combinatorial (as one searches over all possible subsets of $\{\phi_1, \ldots, \phi_{N}\}$ that are good candidates for declaring abstention). In contrast, we can show that if we further require the clusters to be contiguous---namely, if we focus on clustering functions such that $\phi_n < \phi_{n'}< \phi_{n{''}}$ and $c(n)=c(n'')=k$ imply $c(n')=k$---then a minor modification of the flow payoff in the dynamic programming algorithm presented in Section (ref) produces a solution to the problem with contiguous clusters.
Figure (ref) presents the results of the archetype discovery problem with abstention in the context of Example 1. Once again, we take the number of clusters to be $K=10$ and require these clusters to be contiguous. We set the abstention cost at $\sigma^2=.0010$. Figure (ref) shows that under this parameterization, the researcher admits ignorance for the covariates associated to the largest values of the policy effects. This result is not ex ante obvious: according to Figure (ref), the cluster with the smallest policy effects has the largest within-cluster variance. One simple intuition for our findings has to do with the fact that, despite the large within-cluster variance, the left-most cluster in Figure (ref) has the smallest share of observations and therefore does not have a big effect on the loss function.
Figure (ref) confirms this intuition and presents a more detailed comparison of the oracle solutions with and without abstention. The share of observations that belong to the right-most cluster (for which the researcher admits ignorance) is about 30%. This cluster seems to contain the three right-most clusters that arise when abstention is not allowed. In the problem without abstention, the left-most cluster exhibited a high intra-cluster variance. This variance is considerably reduced in the problem with abstention.
Finally, Figure (ref) reports the effects of varying $\sigma^2$ over the archetype sets. As the parameter $\sigma^2$ increases (and admitting ignorance becomes more costly), the researcher admits ignorance for fewer and fewer values of $\phi$. Note that at $\sigma^2=.0002$ (the smallest value of $\sigma^2$ we consider in the upper graph of the figure), the researcher admits ignorance for both the smallest and the largest values of $\phi$. As we discussed before, these are the groups of observations with the largest within-cluster variance.
We now return to Example 1 in Section (ref) and compare the exact dynamic-programming solution with two popular algorithms. The first alternative algorithm to summarize $\phi$ is a (CART) tree as described in breiman2017classification. The tree partitions the space of covariates by greedy recursive axis-aligned splits and is restricted to \(K\) terminal leaves. Recursive decision trees implemented using CART-type recursive partitioning are now widely used to estimate heterogeneous causal treatment effects in experimental and observational studies (although it has been shown recently by cattaneo2025honest that causal trees constructed via standard CART-type partitioning may exhibit poor convergence properties). In our example, we are assuming $\phi$ is known, and we simply want to compare the archetype sets generated by the trees to those generated by the oracle solution. In our implementation, the computational complexity of implementing the tree is \( O\!\left(Kd|\mathcal X|\log|\mathcal X|\right), \) where \(d\) is the dimension of the covariate vector.
The second algorithm is Lloyd's algorithm lloyd1982least. It targets the scalar clustering problem more directly, but it is a local-search heuristic and does not generally guarantee the global minimizer of (ref). Its computational complexity is \( O\!\left(TK|\mathcal X|\right), \) where \(T\) is the number of iterations until convergence.
For the CART benchmark, each terminal leaf is assigned the \(p\)-weighted average of \(\phi\) over that leaf, and the resulting partition is evaluated using the same oracle loss \(L_{oracle}\) used in Section (ref). Figure (ref) displays the resulting partition. The left panel shows the induced partition in covariate space, while the right panel maps the same assignment onto the sorted values of \(\phi\).
The figure shows the geometric restriction imposed by the tree. The exact archetype sets in Example 1 in Section (ref) inherit the circular geometry of the level sets of \(\phi\). By contrast, the tree approximates these sets by recursive axis-aligned rectangles. When the same assignment is viewed on the sorted values of \(\phi\), it is fragmented rather than contiguous.
We next consider Lloyd's algorithm. Unlike the tree, Lloyd's algorithm is applied directly to the scalar values of \(\phi\). It is therefore closer to the clustering problem characterized by Theorem (ref). Figure (ref) shows the resulting partition in covariate space and on the sorted values of \(\phi\).
The clustering based on Lloyd's algorithm is much closer to the interval structure delivered by the exact dynamic program. This is expected: both procedures operate on the scalar values of \(\phi\), rather than imposing a recursive partition on \(x\)-space. The numerical closeness, however, should not be interpreted as an optimality guarantee. Lloyd's algorithm returns the fixed point reached from its initialization and update path, whereas the dynamic program computes the global minimizer of (ref).
Table (ref) summarizes the comparison. The table reports two exact algorithms and two benchmark procedures. The first exact algorithm is Bruce's dynamic program, which is the baseline implementation used in Example (ref). The second is Ckmeans.1d.dp, a faster exact implementation of one-dimensional \(K\)-means wang2011ckmeans; for the weighted case considered here, we use the corresponding weighted extension song2020efficient. Both exact algorithms solve the same one-dimensional clustering problem and therefore attain the same oracle loss. Their difference is computational rather than statistical or decision-theoretic.
The comparison highlights four points. First, exact one-dimensional clustering is computationally feasible at the scale of the examples considered in this paper. Second, faster exact implementations such as Ckmeans.1d.dp preserve the same oracle loss while reducing runtime substantially. Third, the CART tree has a much larger oracle loss in this example. Fourth, Lloyd's algorithm is much closer to the exact objective, but its performance remains that of a heuristic benchmark rather than an exact characterization of the oracle archetype report.
In order to illustrate the usefulness of our theoretical results, we revisit the application discussed in Victor25: an experiment with the government of Haryana in North India designed to analyze the effects of a policy bundle that provided different incentives for immunization across several villages. The goal of the intervention was to increase the takeup of immunization services. In particular, the outcome of interest discussed in chernozhukov2025reply is the number of children (15 months or younger in a given month in a given village) that completed all the vaccines in the immunization schedule.
In their data, there are 25 villages that were randomly assigned to the policy bundle, and 78 control villages. Almost all the treatment and control villages are followed for 12 months (the duration of the intervention).\footnote{See p. 1150 and 1151 of Victor25 for specific details on the treatment.} The total number of village-months observations is 843. Their replication package has 42 covariates available, 39 of which are baseline-village characteristics (such as religion, caste, financial status, baseline immunization, etc). The remaining variables include fixed effects, cluster indicators, and external weights for each village.\footnote{Based on our reading of chernozhukov2025reply and their replication package: fixed effect index district-by-year-month cells, capturing the local time and district context of each observation; the cluster indicator identifies the village, since villages appear repeatedly across months; finally, external weight measures village size and is used as the population weight in the original weighted analysis.}
The first two graphs below report the sorted values of the estimators---or machine learning (ML) proxies---of the conditional average treatment effects using both an Elastic Net and a Neural Network. The sorted values reported in Figure (ref) are generated based on a random $67\%/33\%$-split of the 843 observations. The smaller sample in the split---henceforth, the main data---has 281 observations. The ML proxies, $\widehat{\phi}$, are estimated on the dataset containing $843-281=562$ observations, with 78 villages in the control group and 25 villages in the treatment group. The function $\widehat{\phi}$ is then evaluated at the covariates that appear in the main data. This gives a total of $N=96$ different policy effects, ranging from $-13.803$ to $43.60$ for the Elastic Net estimator, and from $-13.495$ to $23.032$ for the Neural Network estimator. The solid horizontal red-line in Figure (ref) represents the median effect of the policy bundle: $2.814$ for the Elastic Net, and 2.441 for the Neural Net. These two numbers are taken from Table III in Victor25, where the reported median is taken over 250 random data splits. The $K=5$ groups that we generate using the dynamic programming algorithm in Section (ref) are reported in different colors along with their corresponding summary values. For comparison, we also added black dashed lines to represent the quintiles of the policy effects. Computing the ML proxies and running the weighted $K$-means clustering algorithm takes about 1.965 seconds in a personal laptop. Computing the groups based on quantiles takes about 1.793 seconds. This means that the extra computational burden associated to constructing archetype sets using the weighted $K$-means algorithm is minimal.
Figure (ref) shows that the archetype sets generated by weighted $K$-means clustering are markedly different to the groups generated by equally-spaced quantiles of the ML proxies on the specific $67\%/33\%$ data split we consider.\footnote{The weights we use for each value of the proxy share main-sample observations that map to that value. If two village-level observations have the same estimated ML proxy but different village populations, they contribute equally to the weights. } A natural question to ask is whether the difference in the size of the archetype sets will also be present in other random data splits. In order to answer this question, we consider 250 random $67\%/33\%$ data splits and, for each of them, we compute the share of observations that belong to each archetype set (the share is always relative to the smaller dataset that contains about $33\%$ of the observations). The orange dots in Figure (ref) report the median share of observations contained in each of the $K=5$ groups.
Figure (ref) confirms the pattern observed in Figure (ref). The group in the center (Group 3) tends to be considerably larger than the size of the group based on the equally-spaced quantiles (which, by construction, contains $20\%$ of the observations). The groups with the smallest and largest values of the proxies (Groups 1 and 5) tend to be smaller when the groups are constructed using weighted $K$-means. The vertical orange lines in Figure (ref) represent the range of shares observed in 95% of the 250 random data splits (where the data splits corresponding to the lowest 2.5% values and the highest 97.5% values have been excluded from consideration). The black crosses in the figure represent the share of observations corresponding to the data split used to construct Figure (ref).
Since the main object of interest in the archetype discovery problem is the summary of the heterogeneous effects for each archetype set, Figure (ref) reports the median value of $\bar{\phi}$---the summary of the ML proxies for each group. Once again, the median we report is based on the 250 random data splits that were used in Figure (ref).
The orange dots in Figure (ref) represent the median value of the researcher's report when the groups are constructed using weighted $K$-means clustering. The blue dots represent the median value of the report when the groups are constructed based on the quintiles of the ML proxy as recommended by Victor25. The vertical lines represent the range of values of $\bar{\phi}$ observed in 95% of the 250 random data splits (where the lowest 2.5% values and the highest 97.5% have been excluded from consideration).
We note that the values of $\bar{\phi}$ based on weighted $K$-means tend to be more extreme---and noisier---for the groups with the smallest and largest values of the ML proxies. Thus, intuition suggests that if we allow the researcher to admit ignorance---as discussed in Section (ref)---the researcher will likely do so for the groups associated to the smallest and largest values of the ML proxies. Figure (ref) confirms this intuition. In particular, the figure reports the archetype sets for different values of the abstention cost ($\sigma^2$) using the same $67\%/33\%$-data split that was used to construct Figure (ref). The largest value of $\sigma^2$ in the $y$-axis corresponds to the smallest value of $\sigma^2$ for which the researcher never admits ignorance. As discussed in Section (ref), if we focus on the case where the abstention function respects the archetype sets, then we can use a dynamic programming algorithm to solve the oracle archetype discovery problem with abstention. Figure (ref) shows that as we decrease the abstention cost, the researcher starts to admit ignorance. More importantly, the first groups affected are those corresponding to the smallest and largest values of the proxies. This pattern arises with both ML proxies.
Finally, we provide a simple illustration of how the archetype sets constructed via weighted $K$-means clustering can be used to complement the Classification Analysis (CLAN) of Victor25. As noted by Victor25, when the GATES analysis reveals substantial heterogeneity, “it is interesting to know the properties of the subpopulations that are the most and least affected”. Using their CLAN methodology, Victor25 find statistical evidence suggesting that “the villages with low levels of pretreatment immunization are the most affected by the incentives.” We focus on two of the baseline covariates that measure pretreatment immunization levels: the fraction of children receiving measles vaccines by 15 months of age, and the fraction of children receiving measles vaccines at credible locations. Figure (ref) reports the values of these covariates for each group across 250 random data splits. For each split, we label the groups according to the value of their corresponding summary $\bar{\phi}$: the first group is always the one with the lowest values of the ML proxies (and hence the lowest summary), and the fifth group has always the highest values. For each data split, we calculate the median value of each covariate in each of the groups. The vertical and horizontal lines in each of the panels in Figure (ref) represent the range of these medians in 95% of the 250 random data splits, where the lowest 2.5% values and the highest 97.5% have been excluded from consideration. The labeled boxes in each panel represent the median value of the medians. We note that the groups constructed by weighted $K$-means clustering differ very clearly in baseline immunization levels. When the groups are constructed using quintiles of the ML proxy, the top groups seem to be very similar in terms of baseline immunization levels (as measured by the two covariates we consider). Thus, the figure further makes the point that the groups constructed using weighted $K$-means clustering will be very different from the ones constructed using equally-spaced quantiles.
We used decision theory to analyze the archetype discovery problem of breza2025generalizability. In this problem, a researcher wants to summarize $N$ heterogeneous policy effects of interest that vary over a discrete set of covariates. The goal is to partition the set of covariates into $K<N$ groups---the archetype sets---and to provide a summary of the policy effects for each group. Since there are different recommendations in the literature regarding how to present and summarize heterogeneous policy effects, decision theory is useful to guide their evaluation.
The main message of this paper is that, under a weighted mean-squared-error criterion, a procedure analogous to the Sorted Group Average Treatment Effects (GATES) of Victor25 solves the archetype discovery problem (in a sense we made precise through our main theoretical results). The key difference is that, in the optimal procedure, the archetype sets are obtained by weighted $K$-means clustering of the $N$ heterogeneous policy effects, instead of relying on $K$ equally-spaced quantiles. An important observation here is that clustering is based on the $N$ scalar policy effects, not on the covariates, as in the recent work of kim2026causal.
A key component of statistical decision theory is the analysis of the risk function of a statistical decision rule, which requires the specification of an action space, a loss function, and a statistical model. As mentioned in chamberlain2007decision, “once one focuses on a risk function, it is natural to think about criteria like average risk and maximum risk that lend themselves to optimization.” In this paper we analyzed both criteria. In terms of average risk, we showed that the procedure that minimizes average risk for a given prior can be obtained by clustering the different values of the posterior mean estimate of the policy effects of interest. This is a very general result, compatible with nonparametric statistical models and general priors. In terms of maximum or worst-case risk, we introduced a statistical model where the researcher observes an estimator of the policy effects of interest. We showed that an approximately minimax procedure in large samples can be obtained by clustering the estimator of the policy effects. For both the average and maximum risk, the $K$ archetype sets can be found using a simple, and well-known, dynamic programming algorithm. We also analyzed the worst-case regret of the approximately minimax procedure and provided an upper bound on the rate at which it converges to zero.
While there are other recent papers in the literature that have suggested the use of $k$-means clustering for summarizing policy effects---see the recent work of kim2026causal---we are not aware of any other paper in the literature showing that clustering the scalar policy effects (and not the covariates) has some decision-theoretic optimality, despite the high number of citations associated to the GATES procedure and the work of Victor25.\footnote{The 2018 working paper version of chernozhukov2018generic has more than 800 citations in google scholar.}
We think there are several areas for future work. First, throughout this paper, we took the desired number of archetypes as given. It would be interesting to consider the possibility of using data-driven strategies for choosing the number of archetypes. See also Remark 3.1 in kim2026causal. Second, the recent work of nath2026inference has developed an inferential framework for constructing confidence sets for cluster labels. It would be interesting to use such a framework to think about how to conduct inference about archetype sets or to adopt the randomization inference framework of imai2025comment,imai2025statistical. More generally, it would be interesting to think about the best way to communicate uncertainty regarding archetype discovery to an audience of policymakers; perhaps building on the recent results in andrews2024communicating and the framework for scientific communication of andrews2021model. Third, as argued by breza2025generalizability, allowing researchers to admit ignorance and developing tools that identify groups for which more data could be collected can improve the communication of scientific discoveries in social sciences. While we were able to show that, under some assumptions, a dynamic programming algorithm for clustering could be used to construct basins of ignorance, it would be desirable to develop more general results for the archetype discovery problem with abstention. A potentially fruitful connection could be based on the work of DupinNielsen2023PartialKMeans. Instead of allowing the researcher to pay a cost to admit ignorance about the policy effects associated to some covariate value, one could give the researcher a fixed budget in which ignorance can be admitted for $x\%$ of the observations in the sample. The results in DupinNielsen2023PartialKMeans suggest that a problem like this could be computationally tractable, using an algorithm similar in spirit to the one used in this paper.
\onehalfspacing
\doublespacing