Mean Estimation: Median of Means. The Arithmetic Mean is boring... and... | by Jakob Poerschmann | Towards Data Science

If we break it down, machine learning comes back to making statistical inferences. Inference means making statements about an unknown distribution based on a sample generated by that same distribution. Consequently, mean estimation is one of the underlying problems of statistical inference and machine learning. We study samples, drawn from some unknown distribution aiming to understand where the original population is centered. Once that is uncovered, we can move on and learn more about the data generating process.
The Empirical Mean applied to the sample is the most common and straightforward estimator for the population mean. It convinces through its ease of use.
For most students, the Arithmetic Mean must be the first equation they see in their educational careers. So how can this estimator be wrong? Well, it depends.
Let’s look at one simple example. We generate a population generated by a standard Gaussian distribution. We then draw a random sample of 100 data points from the same distribution and compute the Arithmetic Mean as an estimator of the population Mean. Unsurprisingly, our estimation is close to perfect. We denote a squared error of 0.0004.
import numpy as npimport randomnp.random.seed(1813)population = np.random.normal(loc=0, size=1000, scale=1)sample = random.choices(population, k=100)print(np.mean(population))print(np.mean(sample))print((np.mean(population)-np.mean(sample))**2)
So let us move on and always use the Arithmetic Mean, right? We tested and validated its performance. And that makes sense since everyone is using it, right? ... Of course not! This good performance is just due to the wonderfully clean distribution the data has been generated from. But the reality is far from Gaussian. Let’s look at a case that is just a little bit more complex.
population = np.random.standard_t(df=1, size=1000)sample = random.choices(population, k=100)print(np.mean(population))print(np.mean(sample))print((np.mean(population)-np.mean(sample))**2
Things look very different now. With a Student-t distribution, we denote an error squared of 5.806. The thick tails of the sampled Student-t distribution may lead to a large number of extreme values occurring in the sample. This compromises estimation performance. So how do we solve this?
The Median of Means (MoM) estimator is one approach to solve the problem described above. Given a data sample, the MoM estimator shuffles the data points and then splits them into k groups of m data points. It then computes the Arithmetic Mean of each group. Finally, we calculate the median over the resulting k Arithmetic Means.
This technique diminishes the effect that outliers have on the final estimation by splitting the data and only considering the median of the resulting sub-estimations. Let’s see how this improved estimator performs on the previous example.
We can reduce the squared error on the same example to 0.730. A significant improvement to the 5.806 we have seen with the Arithmetic Mean.
However, every technique has its weaknesses. The MoM is dependent on the way we split the data initially. Imagine a case in which all the extremes in the data are sorted into one single sub-group by chance. While unlikely, as data scientists we never want to be dependent on pure chance to such a high degree.
You can find the from-scratch implementation in my repo.
As we know, expected values are every statistician’s friend. Thus let us treat the MoM as a random variable dependent on the specific permutation of the data sample. There is still one problem with this approach. To precisely calculate this estimator, we would need to compute the MoM for every possible sample permutation. That would be n! MoM to compute. Not sexy enough yet!
To make it simple, we can estimate the actual permutation invariant MoM by something that comes close to it. For this, we can fix some number N, as the number of permutations we are going to sample randomly.
Testing this estimator against our running example yields an error squared of 0.332. Success! We were able to improve our estimation even further. It’s time to move out of the kindergarten and become a little bolder in terms of experimentation.
def _rob_median_of_means(seq, n):    res = [_median_of_means(seq, 10) for _ in range(n)]    return np.mean(res)print(_rob_median_of_means(population, 10))print(np.mean(population))print((np.mean(population)-_rob_median_of_means(population, 10))**2)
You can find the from-scratch implementation in my repo.
While we got a basic grasp, we need to understand the behavior of these estimators even better. Therefore let us observe their performance ni dependence on various sample sizes and data generating processes. The plots display the test performance of the introduced estimators. The performance is measured given a large variety of sample sizes and distribution shapes. The Arithmetic Mean is referred to as EM. The MoM and permutation invariant MoM are referred to as MoM and RMoM. Furthermore, green curves display the estimator performance on a standard gaussian and the red curve denote the performance on a Student-t data generating process.
Find the from-scratch implementation of this experiment in my repo.
The results yield some satisfying insights. As expected, the (R)MoM outperforms the Arithmetic Mean for thick-tailed distributions. It is not surprising to observe that the Arithmetic Mean is tough to beat for close to Gaussian distributions. However, we can confidently state that (R)MoM can keep up with Arithmetic Mean estimation performance.
In summary, mean estimation is a challenge that we can solve very well in a laboratory environment. However, once the distributions are no longer normal the challenge might grow increasingly hard. Especially outliers might compromise the performance of the most common mean estimator significantly. The Median of Means represents one of many alternatives to this problem. While the MoM comes with its own flaws, it does improve estimation performance in most cases that are not normal. Also, the computational complexity is controllable over the number of data permutations that we want to include. Thus, why don’t you give it a try in your next model?