Generalizing Koopman-Pitman-Darmois

by johnswentworth15 min read15th Jul 20216 comments

27

Ω 15

RationalityWorld Modeling
Frontpage
Crossposted from the AI Alignment Forum. May contain more technical jargon than usual.

Suppose we want to estimate the mean and variance of a normal distribution from some samples  from the distribution. It turns out that we don’t need to keep around all of the data: only the sample mean  and sample variance  are needed for the estimation problem. From those, we can calculate the entire posterior distribution of . The pair  is a “sufficient statistic” for : a statistic which summarizes all the estimation-relevant information from the data.

In principle, every estimation problem has a sufficient statistic: just let the summary be all the data. But this isn’t very useful. We really want to compress the data, pull out the actually-useful parts and throw away the rest. In the normal distribution example, our  samples of data are each 1-dimensional for a total of  dimensions of data, but the summary is only 2-dimensional, no matter how large  grows.

When will an estimation problem have a fixed-dimensional sufficient statistic as the number of data points increases, like the normal-distribution estimation problem? The Koopman-Pitman-Darmois (KPD) Theorem says that, when estimating distribution parameters from IID samples, there is a fixed-dimensional sufficient statistic if-and-only-if the distribution can be expressed in exponential family form:

… for some vector-valued functions  and positive-real-valued function . The sufficient statistic is then . For instance, for a normal distribution, the parameters  are , and we can use  and .

Intuitively, this says the “coupling” between  and  is only via the fixed-dimensional “summary” vectors  and .

That’s the original KPD Theorem. We’ll see in this post that the theorem generalizes a lot. It’s not just estimation of distribution parameters from IID samples - KPD can also generalize to estimation of parameters for Markov chains, or causal models with symmetry - aka stochastic programs. The theorem also extends to causal models without any symmetry, albeit slightly weaker: a small number of variables can violate the exponential form, but (roughly speaking) the number of variables which can violate the exponential form is proportional to the summary dimension. So if we have a low-dimensional summary of lots of variables, the vast majority of variables must still satisfy the exponential form. The condition which does most of the work is sparsity: we’ll assume the conditional distribution factors according to a causal model with a sparse graph. Or, in more physical terms: direct interactions are local; long-range interactions are always mediated by short-range interactions. Conditionally independent variables are just the most extreme version of this - their conditional distribution factors according to a graph with no edges at all (i.e. there are no interactions at all).

In particular, I now expect most abstractions to mostly satisfy the exponential family form, though that’s a topic for future posts.

The Theorems

We’ll cover three generalized KPD theorems, each with different assumptions on the structure of :

  • Independent But Non-Identical
  • Causal Model/Bayes Net
  • IID

The first is intended as a relatively-simple case to illustrate the core ideas, before we introduce more complexity. The second is the most general theorem. The third is intended to illustrate how we can leverage symmetry to rule out the annoying “exception” variables.

This is definitely not a comprehensive set of KPD generalizations; these theorems and proofs are mainly intended to illustrate the general tools and techniques by which KPD generalizations can be derived.

Independent But Non-Identical KPD

The simplest starting point for these theorems is parameter estimation from independent-but-not-identically-distributed variables. (If you want a concrete image in mind, picture scientists trying to estimate the values of some physical constants using a bunch of different experiments.)

Theorem

Let  be continuous random variables, conditionally independent given  (i.e. ). There exists a D-dimensional sufficient statistic  summarizing all information in  relevant to  only if the distribution has the form

… where:

  •  are some at-most D-dimensional vector-valued summary functions
  •  is a set of “exception” variables, not satisfying the exponential form
  • The number of exception variables  is at-most D

Intuitively, this is like the original KPD, but some variables can be “exceptions” and not follow the exponential form - their coupling with  is not mediated by the summary vectors  and . However, there can only be at-most D exceptions. When the number of variables is much greater than the summary dimension (i.e.  D), this means that the vast majority of variables must satisfy the exponential form. The summary vectors mediate the interactions with  of all but at-most D variables.

Note that this is an “only if” theorem, not an “if-and-only-if” theorem. It does establish that a fixed-dimension sufficient statistic exists if-and-only-if the distribution takes the above form, since any distribution with the above form can use  as a sufficient statistic. However, that sufficient statistic has dimension at-most 2D, whereas the “only if” theorem above assumes the existence of a statistic of dimension D. So there’s a “gap” in the if-and-only-if characterization: any conditionally-independent distribution with a D-dimensional sufficient statistic satisfies the form, but given a distribution which satisfies the form, we may only be able to find a 2D-dimensional sufficient statistic rather than a D-dimensional sufficient statistic.

Of course, for applications with  D, the difference between D and 2D won't matter much anyway.

Proof

The first key idea is the Minimal Map Theorem (also proved more elegantly here): if we want to summarize all the information in  which is relevant to , then the posterior distribution  is a sufficient statistic. It’s also “minimal” in the sense that the posterior distribution can be calculated from any other sufficient statistic, therefore any other sufficient statistic must contain at least as much information (this can be proven via the Data Processing Inequality).

So, given our hypothetical fixed-dimensional summary statistic  (where  is a vector of all the data points), we can write

… for some function . The right-hand side is a data structure containing the values of  for EVERY -value; we can think of it as a vector indexed by . (In the case where  has  possible values , any map  with real-valued  can be represented as a length- vector  in which  is .)

Now, the key property underlying (this version of) the KPD is conditional independence of , i.e. . We want to leverage that factorization, so we’ll transform the distribution into a representation we can factor. First, transform the distribution to a log-odds representation:

… for some arbitrary reference parameters . Next, we apply Bayes’ Rule:

… then subtract off the prior  from both sides, since it’s not a function of :

To clean up the notation, we’ll define . Finally, we can apply the factorization (remember, that was the point of all these transformations) and get

The right-hand side of this expression is a sum of vectors indexed by . Each vector is a function of just one of the , and the left-hand side says that we can express the sum in terms of a D-dimensional summary . So, this is an Additive Summary Equation (see appendix below for details). In fact, it’s a particularly simple Additive Summary Equation: each  depends only on , and each  has no neighbors besides itself.

The Additive Summary Equation Theorem then tells us that the equation is solvable (i.e. a D-dimensional summary statistic exists) only if

… for some  consisting of at-most D -indices, and some  of dimension at-most D.

Exponentiating both sides and simplifying a bit, this condition becomes

… which is the desired form.

Key Takeaways

The main takeaway from the theorem is the idea of “exception” variables, which circumvent the exponential family form, and the idea that the number of exception variables is limited by the dimensionality of the sufficient statistic. 

The main takeaways from the proof are:

  • Start with the Minimal Map Theorem to get the “most general” sufficient statistic
  • Transform to apply the factorization of the distribution
  • Use the Additive Summary Equation Theorem

We’ll reuse this general structure in the next proof.

Causal Model/Bayes Net KPD

Now we address a much more general case: factorization according to a Bayes Net/Causal Model. We still imagine estimating some parameters , but our data is no longer conditionally independent. Conditional on the parameters, interactions between data points are no longer non-existent. However, the direct interactions are sparse - i.e. most data points interact directly with only a few other data points (though they may indirectly interact with everything). This would be the case if, for example, we are trying to estimate parameters of some complicated physical system in real-time from a single instance of the system.

Theorem

Let  be continuous random variables, whose distribution conditional on  forms a Bayes Net on a DAG (i.e. , where  denotes the nodes with edges directly into  in the DAG). There exists a D-dimensional sufficient statistic  summarizing all information in  relevant to  only if the distribution has the form

… where:

  •  are some at-most D-dimensional summary functions
  •  is a set of “exception” variables, not satisfying the exponential form
  • The set of exception variables includes at-most D variables, plus variables in the Markov blanket (i.e. parents, children and parents-of-children) of those D variables, plus children of variables in the Markov blanket.

This mostly generalizes the previous theorem in the obvious way, with one major exception: we can now have a lot more exception variables. Roughly speaking, exceptions happen in “chunks”: neighborhoods each consisting of a variable, its “neighbors” in the Markov blanket, and its neighbors' children. So, in order for the theorem to say anything nontrivial, the DAG has to be sparse: each variable has to have only a few neighbors. If a single variable directly touches most of the DAG, then an exception-neighborhood around that variable could include most of the variables in the model.

As long as the DAG is sparse and  D, the vast majority of variables must still satisfy the exponential form. If each variable has at-most  variables in its Markov blanket, then the number of exception variables is less than D. For instance, in a minimal computationally-complete circuit model (e.g. NAND gates),  (two inputs per node, two children per node to allow copying, and one other parent for each child), so the number of exceptions would be less than 36D. If the number of variables exceeds the summary dimension by several orders of magnitude, this would be a drop in the bucket - the vast majority of variables would still need to satisfy the exponential form.

Proof

This follows broadly the same structure as the proof for the Independent But Non-Identical KPD, so we’ll speed through the parts which are the same and just focus on differences.

As in that proof, we apply the Minimal Map Theorem, then apply a log odds transformation, and factor the distribution. The factorization is now , so the summary equation is

The neighbors  of a variable  are no longer trivial; they are exactly the Markov blanket of  (i.e. parents, children, and children of parents of  in the Bayes net). The terms which depend on  are the factors corresponding to neighbors plus children-of-neighbors, i.e. . The Additive Summary Equation Theorem then says that a D-dimensional summary exists only if

… where , and  is at-most D.

Exponentiating and simplifying a bit yields:

… which is the desired form.

Key Takeaways

The main takeaway here is the idea that exceptional variables occur in neighborhoods, so sparsity is a necessary condition for the theorem to say anything nontrivial. On the other hand, causal models are extremely general, sparsity is ubiquitous in the physical sciences, and a sparse causal model is all we need to apply the theorem. We should thus expect exponential family distributions to show up in most places where low-dimensional summaries exist, in practice. In particular, we should expect exponential family forms for most abstractions, in practice.

Another important point: the methods used in the proof are “smooth”, in the sense that small violations of the assumptions should produce only small violations of the conclusions. In other words, we should also expect distributions with approximate low-dimensional summaries, which approximately factor as a sparse Bayes Net, to approximately fit the exponential form.

IID KPD

When our variables have some symmetry, we can sometimes leverage it to eliminate the annoying “exception” variables.

The basic requirement is that our distribution  is invariant under permuting some of the variables. For instance, if we have a time-symmetric Markov Chain, then we can replace every  with , and the distribution remains the same. In the case of IID variables, we can swap any two variables  and , and the distribution remains the same, i.e.

If we can swap exception variables with non-exception variables, then we can sometimes eliminate the exceptions altogether.

Theorem

Let  be continuous random variables, conditionally independent and identically distributed given  (i.e. ). There exists a D-dimensional sufficient statistic  summarizing all information in  relevant to , with D , if-and-only-if the distribution has the form

… where  are some at-most D-dimensional summary functions.

Other than eliminating the exception terms, there are two differences from the independent-but-not-identically-distributed theorem: the ’s are all the same function, and we explicitly require D . When D  in the earlier theorem, all variables can be exceptions, so the theorem says nothing. For this theorem, we need at least one non-exceptional variable to swap the other variables with, so we require D .

Note that this theorem is strictly stronger than the original KPD as typically stated, which also applies to the IID case. The original KPD talks about existence of a finite-dimensional summary for an infinite stream of IID variables, whereas this version applies even to a finite set of variables, so long as the dimension of the sufficient statistic is smaller than the number of variables.

Proof

Since this a special case of independent variables, we can start from our independent but non-identical theorem:

Since we assume D , there must be at least one non-exception variable. So, without loss of generality, assume .

For any other variable , we can swap  with  and integrate out all the other variables to find

 

… then integrate out :

(Minor note: we’re absorbing constants into  from each integral, which is why it keeps gaining primes.)

For variables , we can also swap with  in order to replace  with . Swapping and integrating as before yields

Substituting both of those into the original expression from the independent but non-identical theorem yields

… which is the desired form.

Key Takeaways

While this is a minor improvement over the original KPD, the real point is to show how symmetry can remove the exception terms.

I expect this to be an especially powerful tool in conjunction with the techniques in Writing Causal Models Like We Write Programs. That post illustrates how to turn recursively-defined functions into recursively-defined causal models. The recursive calls become exactly the sort of symmetries which could potentially be used to remove exception terms in the Causal Model/Bayes Net KPD.

Takeaways Summary

The original Koopman-Pitman-Darmois theorem says that a fixed-dimensional sufficient statistic exists for a stream of IID variables if-and-only-if the distribution is of exponential form.

This generalizes a lot. Neither independence nor identical distribution is necessary; we need only a sparse causal model. However, this generalization comes with a cost: some “exception” variables may not follow the exponential form. As long as the system is sparse, and the number of variables is much larger than the dimension of the sufficient statistic, the number of exception variables will be relatively small, and the vast majority of variables in the system will satisfy the exponential form.

The proofs should also generalize to approximations - systems with approximate sufficient statistics should be approximately exponential family, except for a few variables.

We can also sometimes leverage symmetry to eliminate the exception variables. If the distribution is invariant under some permutation of the variables, then we can try to permute exception variables to non-exception variables, in which case they must follow the exponential form.

Appendix: Additive Summary Equation Theorem

An Additive Summary Equation is a functional equation of the form

… with  and  unknown,  of dimension D , D .

In order to say anything nontrivial about the equation, we need each -term  to only depend on a few -components , and each -component to only influence a few -terms. Given some , we can pick out the (indices of) -components which it influences via the expression

We’re also interested in the “neighbors” of , i.e. -components  for which some  depends on both  and . We’ll write the (indices of) ’s neighbors as ; note that we do consider  a neighbor of itself.

The post on the Additive Summary Equation shows that the equation is solvable only if we can choose a set of at-most D components of , called , for which  can be expressed as

… where the ’s have output dimension at-most D,  is a constant  by D matrix, and  is a constant vector. In particular, we can choose

,

… for some , with  ranging over terms not dependent on , i.e. .

Intuitively, this says that we can’t really say anything about terms dependent on . However,  consists of at-most D variables plus their neighbors, so in a large sparse system, hopefully most terms will not depend on . And for all the terms not dependent on , the information content can be aggregated by summation - i.e. the sum .

Note that we may sometimes want to think about  whose output dimension is uncountably infinite - i.e.  returns a distribution over a continuous variable . In that case, the dot product  defining  becomes an integral ; everything else remains the same. For simplicity, this post's notation implicitly assumes that  has finitely many values, but the generalization is always straightforward.

27

Ω 15

6 comments, sorted by Highlighting new comments since Today at 8:53 AM
New Comment

You've left out the condition for the theorem to be true that the region with non-zero probabilitiy density can't vary with the parameter.

That's how, for example, the uniform distribution over (0,theta) can have a scalar sufficient statistic (the maximum data point) desipite it not being an exponential family model.

Good catch, thanks.

Is there a generalization of "sufficient statistic" that applies to summaries which grow as the log (or polynomial of the log) in the size of the data?

Despite the fact that your summary keeps growing, this seems like it might also be a useful definition. In the limit of infinite data, a log(n) summary is vanishingly small relative to the size of the data you have.

Good question.

The theorems here should apply directly to that situation. The summary will eventually be lower-dimensional than the data, and that's all we need for these theorems to apply. At any data size n sufficiently large that the O(log n) summary dimension is smaller than the data dimension, the distribution of those n data points must be exponential family.

I see! To reiterate: for a fixed n, having a smaller summary dimension means that the distribution of those n points is part of the exponential family.

It seems like there are three cases here:

  1. The size of the summary remains constant as number n of samples grows. The distribution of these samples remains constant. Once enough data is collected, you can estimate the summary statistics and infer the overall distribution.

  2. The size of the summary grows as O(log(n)). The distribution of these samples changes for each n as the size of the summary grows. You converge on the correct distribution in the limit of infinite n. (This seems weird/incorrect, I might be missing something here. I am trying to think of a concrete example)

  3. degenerate case: the size of the required "summary" grows faster than O(n) so you are better off just using the data itself as the summary.

For the second case, each data point can measure something different, possibly correlated with each other, and related in different ways to the parameters we're trying to estimate. For instance, maybe we're trying to estimate some parameters of a car, so we measure the wheel sizes, axle length, number of gears, engine cylinder volume, etc, etc. Every now and then we measure something which gives us a totally different "kind" of information from the other things we measured - something which forces a non-exponential-family update. When that happens, we have to add a new summary component. Eventually other data points may measure the same "kind" of information and also contribute to that component of the summary. But over time, it becomes more and more rare to measure something which no other data point has captured before, so we add summary dimensions more and more slowly.

(Side note: there's no reason I know why O(log n) growth would be special here; the qualitative story would be similar for any sub-linear summary growth.)