Network Flow Methods for the Minimum Covariates Imbalance Problem

The problem of balancing covariates arises in observational studies where one is given a group of control samples and another group, disjoint from the control group, of treatment samples. Each sample, in either group, has several observed nominal covariates. The values, or categories, of each covariate partition the treatment and control samples to a number of subsets referred to as levels where the samples at every level share the same covariate value.

When estimating causal effects using observational data, it is desirable to replicate a randomized experiment as closely as possible by obtaining treatment and control groups with similar covariate distributions. This goal can often be achieved by choosing well-matched samples of the original treatment and control groups, thereby reducing bias due to the covariates. The matching is typically one-to-one.

The bias due to the covariates can be completely removed if for each matched pair, the two samples belong to the same levels over all covariates. Satisfying this requirement, however, typically results in a very small selection from the treatment and control group, which is not desirable. To address this Rosenbaum et al. [1] introduced a relaxation requiring instead to match all treatment samples to a subset of the control samples so that the number of control and treatment samples in each level of each covariate are the same. This requirement is known in the literature as fine balance.

It is not always feasible to find a selection of the control samples that satisfies the fine balance requirement. To that end several papers considered the goal of minimizing the violation of this requirement, which we refer to as imbalance. Yang et al. [2] considered finding an optimal matching between the treatment samples and a selection of the control samples where that selection minimizes an imbalance objective, to be explained below. Bennett et al. [3] considered directly the minimum imbalance problem that finds a selection which minimizes the violation of the fine balance requirement.

Our focus here is the minimum imbalance problem, also referred to as the min-imbalance problem. The min-imbalance problem is trivial for a single covariate and NP-hard for three or more covariates, as discussed next. For the case of two covariates, we introduce here efficient network flow algorithms that solve the problem in polynomial time.

Let $P$ be the number of covariates to be balanced. To introduce the min-imbalance problem, consider first the case of a single covariate, $P=1$, that partitions the control and treatment groups into, say, $k$ levels each. Let the sizes of the levels in the treatment group be $\ell_{1},...,\ell_{k}$ and the sizes of the levels in the control group be $\ell^{\prime}_{1},...,\ell^{\prime}_{k}$. The min-imbalance problem is to select a subset of the control group, called selection, such that the sizes of the levels in the selection, approximate as closely as possible the sizes $\ell_{1},...,\ell_{k}$. The solution to the single covariate min-imbalance problem is trivial: the optimal selection is any $\ell_{i}$ samples out of the $\ell^{\prime}_{i}$ level $i$ control samples if $\ell^{\prime}_{i}\geq\ell_{i}$; otherwise, select all $\ell^{\prime}_{i}$ level $i$ control samples. The value of the objective function corresponding to this selection is then $\sum_{i=1}^{k}\max\{0,\ell_{i}-\ell^{\prime}_{i}\}$, which is the optimal value for the single covariate min-imbalance problem.

For the case of multiple covariates, covariate $p$ partitions both treatment and control groups into $k_{p}$ levels each. Denote the sizes of the levels in the treatment group under covariate $p$ by $\ell_{p,1},\ell_{p,2},...,\ell_{p,k_{p}}$, and let the partition of the control group under covariate $p$ be $L^{\prime}_{p,1},L^{\prime}_{p,2},...,L^{\prime}_{p,k_{p}}$ of sizes $\ell^{\prime}_{p,1},\ell^{\prime}_{p,2},...,\ell^{\prime}_{p,k_{p}}$. The min-imbalance problem is to find a selection $S$, subset of the control group, that minimizes the objective function: $\sum_{p=1}^{P}\sum_{i=1}^{k_{p}}||S\cap L^{\prime}_{p,i}|-\ell_{p,i}|$.

Let the number of treatment samples be $n$ and the number of control samples be $n^{\prime}$. In applications of the min-imbalance problem it is often the case that the size of selection $S$ is required to be $n$, the size of the treatment group. This is assuming that $n\leq n^{\prime}$. Here we focus the presentation on the case where the size of $S$ is required to be $n$. Yet, all our methods generalize to the min-imbalance problem for any specified size $q$, $q\leq n^{\prime}$, of the selection, as shown in Appendix B.

The min-imbalance problem was recently studied by Bennett et al. [3] where a mixed integer programming (MIP) formulation of the problem was presented, with the size of the selection equal to $n$, the size of the treatment group. They showed that the corresponding linear programming (LP) relaxation yields an integer solution when the number of covariates $P\leq 2$. Consequently, Bennett et al. [3] established that the min-imbalance problem on two covariates is solved in polynomial time via linear programming. For $P=3$ they presented an example where the LP relaxation’s solution is not integral. It is further noted in [3], that the min-imbalance problem for $P\geq 3$ is NP-hard. However, no proof or reference is provided. Among other contributions, we fill this gap with a formal proof that the min-imbalance problem for three or more covariates is NP-hard.

A problem related to the min-imbalance that selects a subset of the control samples so as to achieve a certain type of balancing of the levels is called Optimal Matching with Minimal Deviation from Fine Balance, see Yang et al. [2]. We refer to this problem here for short as $\kappa$-Matching-Balance (MB). The MB problem is to optimize the assignment of each treatment sample to a number, $\kappa$, of control samples, based on a distance measure between each treatment and each control sample. This is subject to a constraint that the selection of the $\kappa n$ control samples is an optimal solution to a problem related to min-imbalance. In other words, MB is defined in two stages: In the first stage the goal is to find a selection $S$ of control samples of size $\kappa n$, where $\kappa$ is an integer between $1$ and $\frac{n^{\prime}}{n}$, and $n$ is the size of the treatment group, so as to minimize the $\kappa$-imbalance, defined as $\sum_{p=1}^{P}\sum_{i=1}^{k_{p}}|{|S\cap L^{\prime}_{p,i}|}-\kappa\ell_{p,i}|$. In the second stage, among all selections of control samples that attain the minimum $\kappa$-imbalance, one chooses the selection that minimizes the distance matching of each treatment sample to exactly $\kappa$ selected control samples. Yang et al. [2] studied this MB problem for a single covariate and proposed two network flow algorithms. We observe here that, for any number of covariates, if the first stage problem of minimizing the $\kappa$-imbalance has a unique solution, in terms of the number of selected samples from each level-intersection (to be discussed in the next section, Section 2) of the control group, then an optimal solution to the second stage, and therefore to the MB problem, is easily attained by solving a network flow problem. As shown in Appendix C the minimum $\kappa$-imbalance is equivalent to the min-imbalance problem, and therefore an optimal solution to a corresponding min-imbalance problem provides an optimal solution to the minimum $\kappa$-imbalance problem. It follows in particular that for a single covariate, the first stage problem of minimizing the $\kappa$-imbalance is trivial and has a unique solution in terms of the number of selected samples from each level of the control group. For two or more covariates, no polynomial running time algorithm is known for this problem. However, for two covariates, any of the algorithms shown here, also solve optimally the minimum $\kappa$-imbalance problem. For any number of covariates, when the minimum $\kappa$-imbalance solution is unique (in terms of level-intersection sizes) and known, MB problem is solved in polynomial time in the second stage (for details see Appendix C). In particular for two covariates, if the solution to the minimum $\kappa$-imbalance problem is unique, then the problem is solved efficiently with network flow techniques.

Our main results here are efficient algorithms for the min-imbalance problem with two covariates. We present an integer programming formulation of the problem, related to that of Bennett et al [3], and show that the constraint matrix is totally unimodular. That implies the linear programming relaxation to the problem has all basic solutions, and in particular the optimal solution, integral. We then show that the two-covariate min-imbalance problem can be solved much more efficiently than as a linear program with network flow techniques. We show how to solve the problem as a minimum cost network flow problem with complexity of $O(n\cdot(n^{\prime}+n\log n))$. We then devise a more efficient algorithm based on a maximum flow formulation of the two-covariate min-imbalance problem that runs in $O(n^{\prime 3/2}\log^{2}n)$ steps. Another contribution here is a proof that for three or more covariates, the min-imbalance problem is NP-hard.

In summary our contributions here are:

• •

  An integer programming formulation of the min-imbalance problem with two covariates that has a totally unimodular constraint matrix.

• •

  A formulation of the two-covariate min-imbalance problem as a minimum cost network flow problem, solved in $O(n\cdot(n^{\prime}+n\log n))$ steps.

• •

  A maximum flow formulation of the two-covariate min-imbalance problem that is solved in $O(n^{\prime 3/2}\log^{2}n)$ steps.

• •

  An NP-hardness proof of the min-imbalance problem with three or more covariates.

Paper overview: Section 2 provides the notation used here. Section 3 describes the integer programming formulation with a totally unimodular constraint matrix for the two-covariate min-imbalance problem. We describe the minimum cost network flow formulation of the two-covariate min-imbalance problem, and the algorithm used to solved it in Section 4. In Section 5 we provide a maximum flow problem, the output of which is converted, in linear time, to an optimal solution to the two-covariate min-imbalance problem. We provide several concluding remarks in Section 6. The NP-hardness proof of the min-imbalance problem with three of more covariates is provided in Appendix A. Appendix B explains our method also applies to the min-imbalance problem for other specified size of the selection.

The goal of the min-imbalance problem is to identify a selection of control samples that will match, at each level of each covariate partition, as closely as possible, the number of treatment samples in that level and the number of selection samples in the same level. We denote by $P$ the number of covariates and the number of levels in the partition induced by covariate $p$ by $k_{p}$, for $p=1,...,P$. Let the levels of the treatment group under covariate $p=1,\ldots,P$ be $L_{p,1},L_{p,2},...,L_{p,k_{p}}$ of sizes $\ell_{p,1},\ell_{p,2},...,\ell_{p,k_{p}}$, and let the levels of the control group under covariate $p$ be $L^{\prime}_{p,1},L^{\prime}_{p,2},...,L^{\prime}_{p,k_{p}}$ of sizes $\ell^{\prime}_{p,1},\ell^{\prime}_{p,2},...,\ell^{\prime}_{p,k_{p}}$. For a selection $S$ of control samples we define the discrepancy at level $i$ under covariate $p$ as,

The discrepancy of a level can be positive or negative. If the discrepancy is positive we refer to it as excess which is defined as $e_{p,i}(S)=\max\{0,dis(S,p,i)\}$, and if negative, we refer to it as deficit $d_{p,i}(S)=\max\{0,-dis(S,p,i)\}$. With this notation the imbalance of a selection $S$, $IM(S)$ is,

Notice that this quantity is identical to the imbalance form presented in the introduction: $IM(S)=\sum_{p=1}^{P}\sum_{i=1}^{k_{p}}||S\cap L^{\prime}_{p,i}|-\ell_{p,i}|$.

We now present the integer programming formulation that was given by Bennett et al. in [3] for the min-imbalance problem. That integer program involves two sets of decision variables:
$z_{j}$: a binary variable equal to $1$ if control sample $j$ is in the selection $S$, and $0$ otherwise, for $j=1,\ldots,n^{\prime}$;
$y_{p,i}=|dis(S,p,i)|=||S\cap L^{\prime}_{p,i}|-\ell_{p,i}|$: the absolute value of discrepancy at level $i$ under covariate $p$, for $p=1,\ldots,P$, and $i=1,\ldots,k_{p}$. We note that this variable can only assume integer values, even though Bennett et al. refer to this as a mixed integer programming formulation.
With these variables the formulation is:

For each pair $p,i$, $p=1,\ldots,P$ and $i=1,\ldots,k_{p}$, constraint (1b) and (1c) ensure that $y_{p,i}$ assumes the absolute value of the difference between the number of selected level $i$ control samples and $\ell_{p,i}$ at an optimal solution. These constraints also ensure that any feasible $y_{p,i}$ is non-negative and therefore a non-negativity constraint is not required for variable $y_{p,i}$. The constraint (1d) specifies that the size of selected subset equals the size of the treatment group.

Bennett et al. [3] proved that any basic solution of the linear programming relaxation of (1) is integral for $P=2$. We provide in Section 3 a stronger result showing that a slightly modified form of this integer programming formulation has a constraint matrix which is totally unimodular for $P=2$. That implies that result that every basic solution is integer, and furthermore the constraint matrix is that of a minimum cost network flow problem. This implies that the problem can be solved significantly more efficiently than as linear programming problem.

An optimal solution to formulation (1) specifies for each control sample whether or not it is in the selection. We observe however that the output to the min-imbalance problem, for any number of covariates, can be presented more compactly in terms of level-intersections. For $P$ covariates the intersection of the level sets $L^{\prime}_{1,i_{1}}\cap L^{\prime}_{2,i_{2}}\cap\ldots\cap L^{\prime}_{P,i_{P}}$, $i_{p}=1,\ldots,k_{p}$, $p=1,\ldots,P$, form a partition of the control group. The number of non-empty sets in this partition is at most $\min\{n^{\prime},\Pi_{p=1}^{P}k_{p}\}$. Instead of specifying which sample belongs to the selection, it is sufficient to determine the number of selected control samples in each level intersection, since the identity of the specific selected samples has no effect on the imbalance. With this discussion, we have a theorem on the presentation of the solution to the min-imbalance problem in terms of the level-intersection sizes.

We will say that the solution to the problem is unique if for any optimal selection $S$, the numbers $s_{i_{1},i_{2},\ldots,i_{P}}=|S\cap L^{\prime}_{1,i_{1}}\cap L^{\prime}_{2,i_{2}}\cap\ldots\cap L^{\prime}_{P,i_{P}}|$ are unique. In order to derive an optimal selection given the optimal level-intersection sizes, one selects any $s_{i_{1},i_{2},\ldots,i_{P}}$ control samples from the intersection $L^{\prime}_{1,i_{1}}\cap L^{\prime}_{2,i_{2}}\cap\ldots\cap L^{\prime}_{P,i_{P}}$ for $i_{p}=1,\ldots,k_{p}$, $p=1,\ldots,P$. Obviously, when $s_{i_{1},i_{2},\ldots,i_{P}}<|L^{\prime}_{1,i_{1}}\cap L^{\prime}_{2,i_{2}}\cap\ldots\cap L^{\prime}_{P,i_{P}}|$ for at least one combination of $i_{1},i_{2},\ldots,i_{P}$, the optimal selection is never unique.

Standard formulations of the minimum cost network flow (MCNF) and the maximum flow problems as well as a discussion of why MCNF constraint matrix is totally unimodular, are given in Appendix D.

In our formulation, instead of using variables $y_{p,i}$, we use variables for excess and deficit. As discussed in Section 2, $||S\cap L^{\prime}_{p,i}|-\ell_{p,i}|=e_{p,i}(S)+d_{p,i}(S)$ for each $p$ and $i$. We let the variable for excess for $p$ and $i$ be $e_{p,i}$ and the variable for deficit be $d_{p,i}$. Note that $y_{p,i}=e_{p,i}+d_{p,i}$ where both $e_{p,i}$ and $d_{p,i}$ are non-negative variables.

For each $p$ and $i$, $|S\cap L^{\prime}_{p,i}|-\ell_{p,i}=e_{p,i}-d_{p,i}\Leftrightarrow|S\cap L^{\prime}_{p,i}|+d_{p,i}-e_{p,i}=\ell_{p,i}$.

In the modified formulation shown below, the constraints (2b) and (2c) for the two covariates are separated to facilitate the identification of the total unimodularity property. Since $L^{\prime}_{1,1},...,L^{\prime}_{1,k_{1}}$ is a partition of the control group, $\sum_{i=1}^{k_{1}}|S\cap L^{\prime}_{1,i}|=|S|$. Also, because $\ell_{1,1},...,\ell_{1,k_{1}}$ are the sizes of the levels partition of the treatment group for the first covariate, it follows that $\sum_{i=1}^{k_{1}}\ell_{1,i}=n$. Therefore, $\sum_{i=1}^{k_{1}}\left(e_{1,i}-d_{1,i}\right)=\sum_{i=1}^{k_{1}}\left(|S\cap L^{\prime}_{1,i}|-\ell_{1,i}\right)=|S|-n$ . So specifying $|S|=n$ is equivalent to constraint (2d) in formulation (2) given below:

Formulation (2) is in fact also a minimum cost network flow (MCNF) formulation. See Appendix D for a generic formulation of MCNF. A generic MCNF formulation has exactly one $1$ and one $-1$ in each column of the constraint matrix. To make formulation (2) have this structure, we multiple all coefficients in constraint (2c) by $-1$, and add a redundant constraint $\sum_{i=1}^{k_{2}}d_{2,i}-\sum_{i=1}^{k_{2}}e_{2,i}=0$, which, like constraint (2d), is also equivalent to $|S|=n$. In the next section, we streamline this formulation to derive a more efficient network flow formulation for this problem.

Here we use the level-intersection sizes as variables, $x_{i_{1},i_{2}}$ for $i_{1}=1,...,k_{1},i_{2}=1,...,k_{2}$. These variables can also be written as $x_{i_{1},i_{2}}=\sum_{j\in L^{\prime}_{1,i_{1}}\cap L^{\prime}_{2,i_{2}}}z_{j}$. With these decision variables we get the following network flow formulation:

The formulation (3) is a minimum cost network flow problem. The corresponding network is shown in Figure 1, where all capacity lower bounds are $0$, and each arc has a cost per unit flow and upper bound associated with it. Nodes of type $(1,i_{1})$ each has supply of $\ell_{1,i_{1}}$ and nodes of type $(2,i_{2})$ each has demand of $\ell_{2,i_{2}}$ In Figure 1, for each $i_{1}$ and $i_{2}$, the flow on the arc between node $(1,i_{1})$ and node $(2,i_{2})$ represents variable $x_{i_{1},i_{2}}$; arc from node $1$ to node $(1,i_{1})$ represents the excess $e_{1,i_{1}}$; arc to node $1$ from any node $(1,i_{1})$ represents the deficit $d_{1,i_{1}}$; arc from node $2$ to any node $(2,i_{2})$ represents the deficit $d_{2,i_{2}}$; arc to node $2$ from any node $(2,i_{2})$ represents the excess $e_{2,i_{2}}$. So the per unit arc cost should be 1 for arcs between node 1 or 2 and any node in $\{(1,1),(1,2),...,(1,k_{1})\}\cup\{(2,1),(2,2),...,(2,k_{2})\}$. All other arcs have cost $0$. It is easy to verify that constraints (3b) are corresponding to the flow balance at nodes $(1,i_{1})$ for all $i_{1}$, constraints (3c) are corresponding to the flow balance at nodes $(2,i_{2})$ for all $i_{2}$. Constraint (3d) is corresponding to the flow balance at node 1, and constraint (3e) is corresponding to the flow balance at node 2.

Here we show a maximum flow (max-flow) formulation (see Appendix D for a generic formulation of max-flow problem) for the min-imbalance problem with 2 covariates. Unlike the previous formulations, the maximum flow solution requires further manipulation in order to derive an optimal solution to the min-imbalance problem with 2 covariate. That max-flow graph is illustrated in Figure (2). The source node $s$ can send at most $\ell_{1,i_{1}}$ units of flow to node $(1,i_{1})$ for each $i_{1}=1,...,k_{1}$, the sink node can get at most $\ell_{2,i_{2}}$ units of flow from node $(2,i_{2})$ for each $i_{2}=1,...,k_{2}$, and there can be a flow from node $(1,i_{1})$ to node $(2,i_{2})$ with amount bounded by $u_{i_{1},i_{2}}$, for $i_{1}=1,...,k_{1},i_{2}=1,...,k_{2}$.

Let the maximum flow value for the max-flow problem presented in Figure (2) be denoted by $f^{*}$, and let $\mathbf{x}^{*}$ be the optimal flow vector, with $x^{*}_{i_{1},i_{2}}$ denoting the flow amount between node $(1,i_{1})$ and node $(2,i_{2})$. It is obvious that $\sum_{i_{1}=1}^{k_{1}}\sum_{i_{2}=1}^{k_{2}}x^{*}_{i_{1},i_{2}}=f^{*}\leq\sum_{i_{1}=1}^{k_{1}}\ell_{1,i_{1}}=n$. That means an initial selection $S^{\prime}$ generated by selecting the prescribed number of control samples as in the optimal max-flow solution, i.e., selecting $x^{*}_{i_{1},i_{2}}$ control samples from $L^{\prime}_{1,i_{1}}\cap L^{\prime}_{2,i_{2}}$ is of size $f^{*}$. In order to get a feasible solution for the min-imbalance problem it is required to select $n-f^{*}$ additional control samples. The selection $S^{\prime}$ has no positive excess, only non-negative deficits with respect to the levels of both covariates. This is because $\sum_{i_{2}=1}^{k_{2}}x^{*}_{i_{1},i_{2}}\leq\ell_{1,i_{1}}$ due to the upper bound of arc from $s$ to $(1,i_{1})$ for each $i_{1}$, and $\sum_{i_{1}=1}^{k_{1}}x^{*}_{i_{1},i_{2}}\leq\ell_{2,i_{2}}$ due to the upper bound of arc from $(2,i_{2})$ to $t$ for each $i_{2}$. To recover an optimal solution for the min-imbalance problem from the initial set $S^{\prime}$, we add up to $n-f^{*}$ unselected control samples, one at a time, each corresponding to a level with positive deficit under either covariate $1$ or $2$. This process is repeated until either $n-f^{*}$ such control samples are found, or until no such control sample exists. In the latter case, to complete the size of the selection, any randomly selected control samples are added. Algorithm 1 is a formal statement of this process of recovering an optimal solution of the min-imbalance problem from the initial selection $S^{\prime}$.

To show that Algorithm 1 provides an optimal solution to the min-imbalance problem, we distinguish two cases of Algorithm 1: (1) $S^{+}=\emptyset$ and (2) $|S^{+}|\geq 1$. In the first case, there is, at each iteration, at least one control sample that belongs to some level with positive deficit. In Theorem 3 we prove that the output $S^{\prime}$ of Algorithm 1 is an optimal solution in this case.

We now address the second case where $|S^{+}|\geq 1$ and $|S^{\prime\prime}|<n$. In this case, the total imbalance of $S^{\prime\prime}$ is, from the arguments in the proof of Theorem 3, $IM(S^{\prime\prime})=2(n-f^{*})$. Each one of the $S^{+}$ samples selected adds 1 unit of excess to each covariate, resulting in the addition of 2 units of excess to the imbalance. Therefore, the total imbalance of the output solution is $2(n-f^{*})+2|S^{+}|$. We next show what the value of $|S^{+}|$ is, and then demonstrate that any feasible selection to the min-imbalance problem has total imbalance of at least $2(n-f^{*})+2|S^{+}|$. This will prove that the output of Algorithm 1, $S^{\prime}$, is an optimal solution to the min-imbalance problem.

It will be useful to consider an equivalent form of Algorithm 1. For each level $i$ of covariate $p$ that has $|S^{\prime}\cap L^{\prime}_{p,i}|<\ell_{p,i}$, we add the largest number possible of available control samples in $L^{\prime}_{p,i}$ so long as the total does not exceed $n$. This number is $\min\{\ell_{p,i}-|S^{\prime}\cap L^{\prime}_{p,i}|,\ell^{\prime}_{p,i}-|S^{\prime}\cap L^{\prime}_{p,i}|\}$. Let $\bar{\ell}_{p,i}=\min\{\ell_{p,i},\ell^{\prime}_{p,i}\}$ then for each $p,i$ that has $|S^{\prime}\cap L^{\prime}_{p,i}|<\ell_{p,i}$ we add $\bar{\ell}_{p,i}-|S^{\prime}\cap L^{\prime}_{p,i}|$ previously unselected control samples to $S^{\prime}$. The outcome of this equivalent procedure is exactly the same as that of Algorithm 1. In the case that $|S^{+}|\geq 1$ there is an insufficient number of control samples to add to $S^{\prime}$ after for all levels the largest number possible has been added. Therefore, at the end of this process, the if step returns that another unselected control sample does not exist, and the total number of control samples of $S^{\prime\prime}$, for each level $i$ of covariate $p$, is $\bar{\ell}_{p,i}$.