SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jane H. Lee
jane.h.lee@yale.edu
Andre Wibisono
andre.wibisono@yale.edu
Manolis Zampetakis
emmanouil.zampetakis@yale.edu
Missing data problems have many manifestations across many scientific fields. A fundamental type of missing data problem arises when samples are truncated, i.e., samples that lie in a subset of the support are not observed. Statistical estimation from truncated samples is a classical problem in statistics which dates back to Galton, Pearson, and Fisher. A recent line of work provides the first efficient estimation algorithms for the parameters of a Gaussian distribution [10] and for linear regression with Gaussian noise [11, 14, 37].
In this paper we generalize these results to log-concave exponential families. We provide an estimation algorithm that shows that extrapolation is possible for a much larger class of distributions while it maintains a polynomial sample and time complexity on average. Our algorithm is based on Projected Stochastic Gradient Descent and is not only applicable in a more general setting but is also simpler than the recent algorithms of [10, 26, 11, 14, 37]. Our work also has interesting implications for learning general log-concave distributions and sampling given only access to truncated data.
In many statistical estimation and inference problems, we have access to only a limited part of the data that would be necessary for the classical statistical methods to work, which motivates the development of statistical methods that are resilient to missing data [29]. Truncation [32, 8] is a fundamental and frequent type of missing data and arises when samples that lie outside a subset of the support are not observed and their count is also not observed. Statistical estimation from truncated samples is the focus of the field of truncated statistics, which was developed since the beginning of the twentieth century starting with the work of Galton [19], Pearson and Lee [34, 35], and Fisher [16]. Truncated statistics is widely applicable in Econometrics and many other theoretical and applied fields [32].
A recent line of work establishes the first sample optimal and computationally efficient methods for fundamental statistical estimation problems from truncated samples [10, 26, 11, 13, 14, 24, 37]. All the aforementioned works though heavily rely on the Gaussianity of the distribution of data or the Gaussianity of the noise in regression problems. Gaussianity is an idealized assumption and the question of generalizing truncated statistics beyond Gaussianity has been explored in many existing works, e.g., [1, 22, 38]. The only results in this regime though are for single dimensional problems and truncations that can be described as intervals.
Figure 1: Visualizing the density and truncation set for 2-dimensional exponential distribution. The plot on the left is the density of  $p(x_{1},x_{2})\propto \exp (-\frac{x_{1}}{5} -\frac{x_{2}}{2})$ . On the right is a contour plot of the truncation set  $[0,2]\times [0,2]$  under  $p(x_{1},x_{2})$ . Note that  $\mathbb{E}[X_1] = 5$  and  $\mathbb{E}[X_2] = 2$ , so the truncation set excludes the mean in one direction and includes it only as a boundary point in the other.
In this work we provide statistically and computationally efficient methods for estimating the parameters of exponential families from truncated samples. Our results generalize the recent work of [10] and are the first to provide an estimation algorithm for this problem for a general class of exponential families and for a general class of truncation biases.
Exponential families are one of the most influential type of distribution classes since they include many fundamental distributions such as normal, exponential, beta, gamma, chi-squared, and Weibull distributions. They were first introduced by Fisher [17] and later generalized by [9, 27, 36]. The estimation of the parameters of exponential families over continuous domains is the subject of many classical and recent results; starting from the work of Fisher [17] until the recent results of [25, 41]. This line of work has also found applications in many areas of statistics including causal inference [40]. Our work contributes in this line of research as well since we show how to estimate exponential families even when we only have access to truncated samples.
For every distribution  $p$ , we define the truncated version  $p^S$  as the distribution with density that satisfies  $p^S(\mathbf{x}) \propto \mathbf{1}\{\mathbf{x} \in S\} \cdot p(\mathbf{x})$ . We consider distributions  $p_\theta$  with support  $\mathcal{X} \subseteq \mathbb{R}^m$  parameterized by a vector of parameters  $\theta \in \mathbb{R}^k$  that can be written in the following form:
$$
p _ {\boldsymbol {\theta}} (\mathbf {x}) \propto h (\mathbf {x}) \cdot \exp \left(\boldsymbol {\theta} ^ {\top} T (\mathbf {x})\right)
$$
where  $h$  and  $T$  are known functions and  $T$  is called the sufficient statistics of the exponential family. Our goal is to estimate the vector of parameters  $\theta$  up to  $\epsilon$  error in  $\ell_2$  distance using only samples from  $p_{\theta}^{S}$  when we have oracle access to  $S$ , i.e., there is an oracle that for every  $\mathbf{x} \in \mathcal{X}$  can answer whether  $\mathbf{x} \in S$  or not.
Without any further assumptions estimation from truncated samples is impossible as shown in [10]. In this paper we use the following assumptions to get our estimation guarantees:
Assumption I: We assume a lower bound and an upper bound on the variance of the sufficient statistics in any direction. This assumption is used in estimation of the parameters of exponential families even without truncated samples, e.g., Assumption 4.1 in [41].
Assumption II: We assume that the exponential families contain log-concave distributions. We use this assumption in two places: (1) to show our extrapolation result, our distributions must satisfy an anti-concentration property; this is a property that is heavily utilized even in estimating a Gaussian distribution [10], (2) we require access to a sampling oracle for the underlying distributions; this sampling oracle can be implemented efficiently using Langevin dynamics if we have log-concavity.
Assumption III: We assume that the sufficient statistics  $T$  are bounded degree polynomials. This assumption combined with log-concavity provides the anti-concentration property that we
need. Alternatively, we can assume that the sufficient statistics belongs to a function class that satisfies anti-concentration but for simplicity of exposition we focus on the case where  $T(\mathbf{x})$  is a polynomial.
Our main result is the following.
Informal Theorem 1.1 (See Theorem 3.1 for the formal statement). Under Assumptions I, II, III, and given samples  $n$  (at least  $\widetilde{\mathcal{O}}(k/\epsilon^2)$  for some  $\epsilon$ ) from  $p_{\pmb{\theta}^*}^S$ , where  $S$  is a measurable set to which we have oracle access, there exists an estimation algorithm with (expected) running time  $\text{poly}(m, k, n)$  that outputs  $\widehat{\pmb{\theta}}$  such that with probability at least  $99\%$ , it holds that  $\| \widehat{\pmb{\theta}} - \pmb{\theta}^* \| < \epsilon$ .
Our main result has the following important implications:
> We show that our assumptions are satisfied from exponential distributions, Weibull distributions, continuous Bernoulli, continuous Poisson, Gaussian distributions, and generalized linear models. Hence, our result implies an efficient method for estimation from truncated samples for all these distribution classes.
$\triangleright$  Another interesting corollary of our result is that we can combine it with the ideas of [12] and get a general method for learning log-concave distributions from truncated samples. In particular, assume that we want to learn from truncated samples a distribution that can be written in the form  $p(\mathbf{x}) \propto \exp(-f(\mathbf{x}))$  where  $f$  is a convex function. Now under mild assumptions, we can replace  $f(\mathbf{x})$  with a finite Taylor approximation, i.e., we have that  $f(\mathbf{x}) \cong \sum_{i} a_{i} t_{i}(\mathbf{x})$  for some polynomials  $t_{i}(\mathbf{x})$ . Then, using our method we can estimate the parameters  $a_{i}$  and output an estimation of  $p$ .
$\triangleright$  In the context of sampling, the gradient of the log-likelihood is the score function that is needed to run Langevin dynamics (e.g., see [30, 7]). Our result also says that we are able to sample from the original distribution, given that we only observe truncated samples.
Technical Contributions. Our algorithm is projected stochastic gradient descent (PSGD) on the negative log-likelihood function and is the same algorithm that is used in many of the recent works in truncated statistics. As in the previous work though, the main challenge is to show that PSGD will converge in our general setting. The previous analysis of the convergence of PSGD was based on exact properties of the Gaussian distribution. When we move away from Gaussianity every step of the convergence analysis becomes much more technical and we need to make sure to only use properties that generalize to exponential families beyond Gaussian distributions. In particular, certain results such as that of Lemma 3.6 and its related quantities in C.1 can be generalized even beyond exponential families, holding for any density.
Our most related literature is the recent series of works on truncated statistics which includes the following results: estimation of multivariate normal distributions [10], linear regression with Gaussian noise [10, 26, 11, 13, 14, 37], estimation of product distributions over the hypercube [18], non-parametric density estimation [12]. All of these works heavily rely on properties of the Gaussian distributions, or product distributions over the hypercube, or their dependence in the number of dimensions is not efficient, e.g., [12]. In our work we identify the properties of exponential families that are only required to get the efficient estimation results and we show that linear dependence on the dimension is achievable in settings that are more general than the Gaussian case.
Another related work is that of [30] that solves parameter estimation of a truncated density given samples through the score matching technique. To derive a tractable objective, we need appropriate boundary conditions which are not satisfied by truncated densities, but [30] instead uses a modified weighted Fisher distance given that the truncation set  $S$  is a Lipschitz domain (a type of open and connected set). On the other hand, our work assumes no particular structure about  $S$  and hence our results are more general and applicable in much more complicated settings for exponential families.
Notation. Lowercase bold letters will denote real-valued vectors, e.g.,  $\mathbf{x} \in \mathbb{R}^m$ , and uppercase bold letters will denote matrices with real values, e.g.,  $\mathbf{A} \in \mathbb{R}^{n \times m}$ . For a random vector  $\mathbf{x} \sim \rho$ ,
$\mathbf{Cov}[\mathbf{x}] = \mathbf{Cov}[\mathbf{x},\mathbf{x}] = \mathbb{E}[(\mathbf{x} - \mathbb{E}[\mathbf{x}])(\mathbf{x} - \mathbb{E}[\mathbf{x}])^\top ]$  is its covariance matrix, and  $\mathbf{Var}(\mathbf{x})$  is the trace of the covariance matrix (a scalar value). Depending on whether it is clear from context, Cov and Var may include subscripts to indicate the distribution  $\rho$ . The notation  $B(c,R)$  is the Euclidean ball centered at  $c$  with radius  $R > 0$ .
Exponential Families. Let  $\mathbf{x} \in \mathcal{X} \subseteq \mathbb{R}^m$ . We are interested in a class of densities which have the form,
$$
p _ {\theta} (\mathbf {x}) = h (\mathbf {x}) \exp (\boldsymbol {\theta} ^ {\top} T (\mathbf {x}) - A (\boldsymbol {\theta})),
$$
where  $h: \mathbb{R}^m \mapsto \mathbb{R}_+$  is the base or carrier measure,  $\pmb{\theta} \in \Theta$  with  $\Theta = \{\pmb{\theta} \in \mathbb{R}^k : A(\pmb{\theta}) < \infty\}$  is the natural parameter space,  $T: \mathbb{R}^m \mapsto \mathbb{R}^k$  is the sufficient statistic for  $\pmb{\theta}$ , and  $A(\pmb{\theta}) = \log Z(\pmb{\theta})$  is the log-partition function, where  $Z(\pmb{\theta}) = \int p_{\pmb{\theta}}(\mathbf{x}) d\mathbf{x}$ .
A regular exponential family is one where  $\Theta$  is an open set. It is minimal if the  $\theta$  and  $T(\mathbf{x})$  are each linearly independent. Any non-minimal family can be made minimal by appropriate reparametrization. In any regular exponential family,  $A(\theta)$  is convex. It is strictly convex if the representation is minimal. Exponential families have several nice properties (e.g., see Theorem 1 of [4]), among which are that  $\nabla A(\theta) = \mathbb{E}_{p_{\theta}}[T(\mathbf{x})]$  and  $\nabla^2 A(\theta) = \mathbf{Cov}_{p_{\theta}}[T(\mathbf{x})]$ .
Truncated Distributions. Let  $\rho$  be a probability distribution on  $\mathbb{R}^m$ . We represent  $\rho$  as a probability density function with respect to the Lebesgue measure  $d\mathbf{x}$  on  $\mathbb{R}^m$ . Let  $S \subseteq \mathbb{R}^m$  be such that  $\rho(S) = \alpha$  for some  $\alpha \in (0,1]$ . Let  $\rho^S \coloneqq \rho(\cdot | \cdot \in S)$  be the conditional distribution of  $\mathbf{x} \sim \rho$  given that  $\mathbf{x} \in S$ . Concretely, the density of  $\rho^S$  is
$$
\rho^ {S} (\mathbf {x}) = \frac {\rho (\mathbf {x}) \cdot \mathbb {1} \left\{\mathbf {x} \in S \right\}}{\rho (S)}.
$$
For exponential families, we have the truncated density  $p_{\theta}^{S}(\mathbf{x})$  is:
$$
p _ {\pmb {\theta}} ^ {S} (\mathbf {x}) = \frac {p _ {\pmb {\theta}} (\mathbf {x})}{\int_ {S} p _ {\pmb {\theta}} (\mathbf {x}) d \mathbf {x}} \mathbb {1} \{\mathbf {x} \in S \} = \frac {h (\mathbf {x}) \exp (\pmb {\theta} ^ {\top} T (\mathbf {x}))}{\int_ {S} h (\mathbf {x}) \exp (\pmb {\theta} ^ {\top} T (\mathbf {x})) d \mathbf {x}} \mathbb {1} \{\mathbf {x} \in S \}.
$$
See Figure 1 for an illustration.
Sub-Exponential Distributions. Although the term sub-exponential has been overloaded (e.g., [21] v.s.[44]), the definition we will use describes a class of distributions whose tails decay at least as fast as an exponential, but with potentially heavier tails than Gaussians [44].
There are several equivalent characterizations of sub-exponential random variables (e.g., see Prop. 2.7.1 of [44]), one of which uses the moment generating function.
Definition (Sub-exponential random variable). A centered, real-valued random variable  $X \in SE(\nu^2, \beta)$  is sub-exponential with parameters  $\nu^2, \beta > 0$  if
$$
\mathbb {E} [ e ^ {\lambda X} ] \leq e ^ {\frac {\nu^ {2} \lambda^ {2}}{2}}, \forall \lambda : | \lambda | <   1 / \beta .
$$
Membership Oracle of a Set. Let  $S \subseteq \mathbb{R}^m$ . A membership oracle is an efficient procedure which computes  $\mathbb{1}\{\mathbf{x} \in S\}$ .
3 Projected Stochastic Gradient Descent Algorithm
Problem Setup. We are given truncated samples  $\{\mathbf{x}_i\}_{i=1}^n$ , with each  $\mathbf{x}_i \sim p_{\theta^*}^S$ , where  $p_{\theta^*}(S) = \alpha > 0$ . Without knowledge of the truncation set  $S$  beyond access to a membership oracle, can one recover  $\theta^*$  and thus  $p_{\theta^*}$  efficiently?
We answer this question positively, under the following assumptions:
Assumption A1 (Strong Convexity, Smoothness of Non-truncated Negative Log-Likelihood over  $\Theta$ ).
$$
\lambda I \preceq \mathbf {C o v} _ {\mathbf {z} \sim p _ {\theta}} [ T (\mathbf {z}), T (\mathbf {z}) ] \preceq L I \quad \forall \boldsymbol {\theta} \in \Theta ,
$$
for some  $\lambda, L > 0$ . Here, we've abused notation for  $\Theta$  which can be a subset of the entire natural parameter space. As mentioned earlier, this is always at least convex for exponential families and strictly convex in minimal representation. Thus the negative log-likelihood (of the non-truncated density) can be made strongly convex and smooth by restricting the natural parameter space appropriately.
Assumption A2 (Log-Concave Density). The density  $p_{\theta}(\mathbf{x})$  is log-concave in  $\mathbf{x}$
Assumption A3 (Sufficient Statistics  $T(\mathbf{x})$  is polynomial in  $\mathbf{x}$ ).  $T(\mathbf{x}) \in \mathbb{R}^k$  has components which are polynomial in  $\mathbf{x}$ , with degree at most  $d$ .
Assumptions A2 and A3 allow us to use the anti-concentration result needed for Lemma 3.2 which is heavily utilized even in the Gaussian case. While A2 also allows for efficient sampling via Langevin dynamics, the latter is only used in Lemma 3.2. Refer back to 1.1 for discussion of these assumptions.
Main Result.
Theorem 3.1 (Main). Given membership oracle access to a measurable set  $S$  whose measure is some constant  $\alpha \in (0,1]$  under an unknown exponential family distribution  $p_{\theta^*}$  which satisfies A1, A2, A3, and given samples  $\mathbf{x}_1, \ldots, \mathbf{x}_n$  from  $p_{\theta^*}$  that are truncated to this set, there exists an expected polynomial-time algorithm that recovers an estimate  $\widehat{\theta}$ . That is, for any  $\epsilon > 0$  the algorithm
- Uses an expected  $\widetilde{\mathcal{O}}(k/\epsilon^2)$  truncated samples and queries to the membership oracle,
- Runs in expected poly  $(m, k, 1 / \epsilon)$  time.
- Produces an estimate  $\widehat{\theta}$  such that with probability at least  $99\%$ ,
$$
\left\| \widehat {\boldsymbol {\theta}} - \boldsymbol {\theta} ^ {*} \right\| <   \epsilon .
$$
In order the solve this problem, we need to define an objective whose optimum is  $\theta^{*}$  and we need to be able to recover it uniquely. To use maximum likelihood estimation (or minimize the negative log-likelihood), we have to be able to compute gradients which depend on the truncation set  $S$ , which we cannot do directly without more knowledge about  $S$ . However, we can sample unbiased estimates of the gradient, as long we have non-trivial mass on  $S$  at a current parameter estimate (otherwise the truncated likelihood function at that parameter is not well-defined and rejection sampling would take infinite time). To address all of these issues, the organization of this section is as follows:
- Section 3.1 establishes that after truncation, the negative log-likelihood remains strongly convex and smooth (in  $\theta$ ) over a subset of parameters which have non-trivial mass on the truncation set.
- In Section 3.3, we show that while we do not know the truncation set, we can solve the non-truncated MLE problem with truncated samples to find an initial parameter  $\theta_0$  which assigns non-trivial mass to the truncation set.
- Then given this  $\theta_0$ , in Section 3.2 we show that we can construct a set of parameters  $K$  which all assign non-trivial mass to the truncation set (and contains the true parameter  $\theta^*$ ).
- In Section 3.4, we use results from the previous sections to prove that we can efficiently recover the true parameter  $\theta^{*}$  using a stochastic gradient descent procedure minimizing the truncated negative log-likelihood, which projects to the parameter space  $K$ .
3.1 Strong Convexity and Smoothness of Truncated Negative Log-Likelihood
Without truncation, recovering the true parameter  $\theta^{*}$  for any parameterized distribution given samples is a classical problem solved by maximizing the likelihood (or minimizing its negation). Here, we state the main objective we will minimize through a stochastic gradient descent procedure as well as the properties of this objective that will allow us to recover  $\theta^{*}$ . Define:
$$
\bar {\ell} (\boldsymbol {\theta}) := - \mathbb {E} _ {\mathbf {x} \sim p _ {\boldsymbol {\theta} ^ {*}} ^ {S}} \left[ \log \left(p _ {\boldsymbol {\theta}} ^ {S} (\mathbf {x})\right) \right]
$$
$$
\nabla_ {\pmb {\theta}} \bar {\ell} (\pmb {\theta}) = \mathbb {E} _ {\mathbf {z} \sim p _ {\pmb {\theta}} ^ {S}} [ T (\mathbf {z}) ] - \mathbb {E} _ {\mathbf {x} \sim p _ {\pmb {\theta} ^ {*}} ^ {S}} [ T (\mathbf {x}) ]
$$
$$
\nabla_ {\pmb {\theta}} ^ {2} \bar {\ell} (\pmb {\theta}) = \mathbf {C o v} _ {\mathbf {z} \sim p _ {\pmb {\theta}} ^ {S}} [ T (\mathbf {z}), T (\mathbf {z}) ]
$$
Note that since the Hessian is a covariance matrix which is at least PSD, this objective is always convex in  $\theta$ . Thus  $\theta^{*}$  is a minimizer since it satisfies the first-order optimality condition. (These calculations can be found in Appendix A.1.) However, if the objective is too flat, we may not be able to recover  $\theta^{*}$  even after sufficiently reducing the objective value. For this, we prove that if the original non-truncated covariance has bounded eigenvalues, the truncated one does as well under A1, A2, and A3 at parameters which assign non-trivial mass to  $S$ .
Lemma 3.2 (Preservation of Strong Convexity under Truncation). Assume the lower bound in A1, A2, A3. If  $p_{\theta}(S) > 0$ , then
$$
\mathbf {C o v} _ {\mathbf {z} \sim p _ {\boldsymbol {\theta}} ^ {S}} [ T (\mathbf {z}), T (\mathbf {z}) ] \succeq \frac {1}{2} \left(\frac {p _ {\boldsymbol {\theta}} (S)}{4 C d}\right) ^ {2 d} \lambda I,
$$
where  $C$  is a universal constant guaranteed by Theorem 8 of [5] and  $d$  is the maximum degree of  $T(\mathbf{x})$ . See proof in Appendix A.2 which follows that of [10].
Lemma 3.3 (Preservation of Smoothness under Truncation). Assume the upper bound in A1. Suppose  $p_{\theta}(S) > 0$ , then
$$
\mathbf {C o v} _ {\mathbf {z} \sim p _ {\theta} ^ {S}} [ T (\mathbf {z}), T (\mathbf {z}) ] \preceq \frac {1}{p _ {\theta} (S)} L I.
$$
See proof in Appendix A.3. The proof is simple and can be done similarly to the previous lemma.
Thus, we have shown that as long as we optimize over a parameter space where every  $\theta$  assigns non-trivial mass to the truncation set, our objective is both strongly convex and smooth. The following sections will help us determine and then construct this set given samples.
Remark. Note that the upper bound increased and the lower bound decreased for the eigenvalues of the truncated covariance matrix. Even in a one-dimensional simple case like  $\mathcal{N}(0,1)$ , it is easy to construct examples of both increasing the shrinking the variance given freedom to place mass anywhere under  $\mathcal{N}(0,1)$ . Thus, it may be natural that the eigenvalue range expands after truncation.
Remark. Lemma 3.2 and the log-likelihood calculations are direct generalizations of prior work in the Gaussian case, where we can recover the results of [10] by noting that the re-parameterization of Gaussian parameters  $(\mu, \Sigma)$  as  $\nu = \Sigma^{-1}\mu$  and  $\mathbf{T} = \Sigma^{-1}$  is the natural parameterization of multivariate Gaussian distributions in exponential family form (up to constants). The sufficient statistics here  $T(\mathbf{x}) \propto [\mathbf{x}, \mathbf{xx}^{\top}]$  has components which are polynomial in  $\mathbf{x}$  with degree at most 2, and plugging in  $d = 2$  to Lemma 3.2 recovers Lemma 4 in [10]. Appendix B includes other examples beyond Gaussians which satisfy A1, A2, and A3.
3.2 Parameter Space with Non-Trivial Mass on Truncation Set
The prior section established that the strong convexity and smoothness parameter of the truncated objective is controlled by the mass that  $p_{\theta}$  assigns to the truncation set  $S$  for any given  $\theta$ . In this section, we will prove lower bounds on the mass that  $p_{\theta}$  assigns to the truncation set, given that the parameter distance to the optimum  $\| \pmb{\theta} - \pmb{\theta}^{*} \|$  is bounded.
Lemma 3.4 (Lower bound for mass on truncation set under smoothness given parameter distance). Assume A1. Let  $\theta, \theta' \in \Theta$ . Then for two distributions from the same exponential family
$$
p _ {\boldsymbol {\theta}} (S) \geq p _ {\boldsymbol {\theta} ^ {\prime}} (S) ^ {2} \cdot \exp \left(- \frac {3 L}{2} \| \boldsymbol {\theta} - \boldsymbol {\theta} ^ {\prime} \| ^ {2}\right).
$$
Proof is provided in Appendix C.3, and only needs smoothness. Thus, we can lower bound the mass that a parameter  $\theta$  assigns to  $S$  given its distance  $\| \pmb {\theta} - \pmb{\theta}^{*}\|$  from  $\pmb{\theta}^{*}$  which is assumed to have  $p_{\theta^*}(S) = \alpha$