Set balancing

Formally, the set balancing problem can be described as follows.

${\displaystyle m}$ is the number of subjects in the general population.

${\displaystyle n}$ is the number of potential features.

The subjects are described by ${\displaystyle A}$, an ${\displaystyle n\times m}$ matrix with entries in ${\displaystyle {0,1}}$. Each column represents a subject and each row represents a feature. ${\displaystyle a_{i,j}=1}$ if subject ${\displaystyle j}$ has feature ${\displaystyle i}$, and ${\displaystyle a_{i,j}=0}$ if subject ${\displaystyle j}$ does not have feature ${\displaystyle i}$.

The partition to groups is described by ${\displaystyle b}$, an ${\displaystyle m\times 1}$ vector with entries in ${\displaystyle {-1,1}}$. ${\displaystyle b_{j}=1}$ if subject ${\displaystyle j}$ is in the treatment group T and ${\displaystyle b_{j}=-1}$ is subject ${\displaystyle j}$ is in the control group C.

The balance of features is described by ${\displaystyle c=A\cdot b}$. This is an ${\displaystyle n\times 1}$ vector. The numeric value of ${\displaystyle c_{i}}$ is the imbalance in feature ${\displaystyle i}$: if ${\displaystyle c_{i}>0}$ then there are more subjects with ${\displaystyle i}$ in T and if ${\displaystyle c_{i}<0}$ then there are more subjects with ${\displaystyle i}$ in C.

The imbalance of a given partition is defined as:

${\displaystyle I(b)=||A\cdot b||_{\infty }=\max _{i\in 1\dots ,n}|c_{i}|}$

The set balancing problem is to find a vector ${\displaystyle b}$ which minimizes the imbalance ${\displaystyle I(b)}$.

An approximate solution can be found with the following very simple randomized algorithm:

Send each subject to the treatment group with probability 1/2.

In matrix formulation:

Choose the elements of ${\displaystyle b}$ randomly with probability 1/2 to each value in {1,-1}.

Surprisingly, although this algorithm completely ignores the matrix ${\displaystyle A}$, it achieves a small imbalance with high probability when there are many features. Formally, for a random vector ${\displaystyle b}$:

${\displaystyle Prob\left[I(b)\geq {\sqrt {4m\ln n}}\right]\leq {\frac {2}{n}}}$

PROOF:

Let ${\displaystyle k_{i}}$ be the total number of subjects that have feature ${\displaystyle i}$ (equivalently, the number of ones in the ${\displaystyle i}$-th of the matrix ${\displaystyle A}$). Consider the following two cases:

Easy case: ${\displaystyle k_{i}\leq {\sqrt {4m\ln n}}}$. Then, with probability 1, the imbalance in feature ${\displaystyle i}$ (that we marked by ${\displaystyle c_{i}}$) is at most ${\displaystyle {\sqrt {4m\ln n}}}$.

Hard case: ${\displaystyle k_{i}>{\sqrt {4m\ln n}}}$. For every ${\displaystyle j}$, let ${\displaystyle X_{j}=a_{i,j}b_{j}}$. Each such ${\displaystyle X_{j}}$ is a random variable that can be either 1 or -1 with probability 1/2. The imbalance in feature ${\displaystyle i}$ is: ${\displaystyle c_{i}=\sum _{j=1}^{m}{X_{j}}}$. Since the ${\displaystyle X_{j}}$ are independent random variables, by the Chernoff bound, for every ${\displaystyle a>0}$:

${\displaystyle Prob\left[|c_{i}|\geq a\right]\leq 2\exp(-a^{2}/2m)}$

Select: ${\displaystyle a={\sqrt {4m\ln n}}}$ and get:

${\displaystyle Prob\left[|c_{i}|\geq {\sqrt {4m\ln n}}\right]\leq {\frac {2}{n^{2}}}}$

By the union bound,

${\displaystyle Prob\left[\exists i:|c_{i}|\geq {\sqrt {4m\ln n}}\right]\leq {\frac {2}{n}}}$.