# End-to-end speaker segmentation for overlap-aware resegmentation

###### Abstract

Speaker segmentation consists in partitioning a conversation between one or
more speakers into speaker turns. Usually addressed as the late combination of
three sub-tasks (voice activity detection, speaker change detection, and
overlapped speech detection), we propose to train an end-to-end segmentation
model that does it directly. Inspired by the original end-to-end neural
speaker diarization approach (EEND), the task is modeled as a multi-label
classification problem using permutation-invariant training. The main
difference is that our model operates on short audio chunks (5 seconds) but at
a much higher temporal resolution (every 16ms). Experiments on multiple
speaker diarization datasets conclude that our model can be used with great
success on both voice activity detection and overlapped speech detection. Our
proposed model can also be used as a post-processing step, to detect and
correctly assign overlapped speech regions. Relative diarization error rate
improvement over the best considered baseline (VBx) reaches 18% on AMI, 17% on
DIHARD 3, and 16% on VoxConverse.

Index Terms: speaker diarization, speaker segmentation, voice activity
detection, overlapped speech detection, resegmentation.

##  1 Introduction

The speech processing community relies on term segmentation to describe a
multitude of tasks: from classifying the audio signal into three classes
{speech, music, other}, to detecting breath groups, localizing word
boundaries, or even partitioning speech regions into phonetic units. On this
coarse-to-fine time scale, speaker segmentation lies somewhere between
{speech, non-speech} classification and breath groups detection. It consists
in partitioning speech regions into smaller chunks containing speech from a
single speaker. It has been addressed in the past as the combination of
several sub-tasks. First, voice activity detection (VAD) removes any region
that does not contain speech. Then, speaker change detection (SCD) partitions
remaining speech regions into speaker turns, by looking for time instants
where a change of speaker occurs [1]. From a distance, this definition of
speaker segmentation may appear clear and unambiguous. However, when looking
more carefully, lots of complex phenomena happen in real-life spontaneous
conversations – overlapped speech, interruptions, and backchannels being the
most prominent ones. Therefore, researchers have started working on the
overlapped speech detection (OSD) task as well [2, 3, 4].

End-to-end speaker segmentation. Instead of addressing voice activity
detection, speaker change detection, and overlapped speech detection as three
different tasks, our first contribution is to train a unique end-to-end
speaker segmentation model whose output encompasses the aforementioned sub-
tasks. This model is directly inspired by recent advances in end-to-end
speaker diarization and, in particular, the growing End-to-End Neural
Diarization (EEND) family of approaches developed by Hitachi [5, 6, 7]. The
proposed segmentation model is better than (or at least on par with) several
voice activity detection baselines, and sets a new state of the art for
overlapped speech detection on all three considered datasets: AMI Mix-Headset
[8], DIHARD 3 [9, 10], and VoxConverse [11]. We did not run speaker change
detection experiments.

Overlap-aware resegmentation. Our second contribution relates to the problem
of assigning detected overlapped speech regions to the right speakers. While a
few attempts have been made in the past [4, 12], it remains a very difficult
problem for which a simple heuristic baseline has yet to be beaten [13]. We
show, through extensive experimentation, that our segmentation model
consistently beats this heuristic when turned into an overlap-aware
resegmentation module – setting a new state of the art on the AMI dataset when
combined with the VBx approach.

Reproducible research. Last but not least, our final contribution consists in
sharing the pretrained model with the community and integrating it into
pyannote open-source library for reproducibility purposes:
huggingface.co/pyannote/segmentation.

##  2 End-to-end speaker segmentation

Like in the original EEND approach [5], the task is modeled as a multi-label
classification problem using permutation-invariant training. As depicted in
Figure 1, the main difference is that our model operates on short audio chunks
(5 seconds) but at a much higher temporal resolution (around every 16ms).
Processing short audio chunks also implies that the number of speakers is much
smaller and less variable than with the original EEND approach (dealing with
whole conversations) – making the problem easier to address. For instance, we
found that $99\%$ of every possible 5s chunks in the training set (later
defined in Section 3) contained less than $K_{\text{max}}=4$ speakers.

Figure 1: Actual outputs of our model on two 5s excerpts from the same
conversation between two speakers (source: file DH_EVAL_0035.flac in DIHARD3
dataset). Top row shows the reference annotation. Middle row is the audio
chunk ingested by the model. Bottom row depicts the raw speaker activations,
as returned by the model. Thanks to permutation-invariant training, notice how
the blue speaker corresponds to the orange activation on the left and to the
green one on the right.

###  2.1 Permutation-invariant training

Given an audio chunk X, its reference segmentation can be encoded into a
sequence of $K_{\text{max}}$-dimensional binary frames
$\mathbf{y}=\\{\mathbf{y_{1}},\ldots,\mathbf{y_{T}}\\}$ where
$\mathbf{y_{t}}\in\\{0,1\\}^{K_{\text{max}}}$ and $y_{t}^{k}=1$ if speaker $k$
is active at frame $t$ and $y_{t}^{k}=0$ otherwise. We may arbitrarily sort
speakers by chronological order of their first activity but any permutation of
the $K_{\text{max}}$ dimensions is a valid representation of the reference
segmentation. Therefore, the binary cross entropy loss function
$\mathcal{L}_{\text{BCE}}$ (classically used for such multi-label
classification problems) has to be turned into a permutation-invariant loss
function $\mathcal{L}$ by running over all possible permutations
$\text{perm}(\mathbf{y})$ of $\mathbf{y}$ over its $K_{\text{max}}$
dimensions:

$\mathcal{L}\left(\mathbf{y},\mathbf{\hat{y}}\right)=\min_{\text{perm}(\mathbf{y})}\mathcal{L}_{\text{BCE}}\left(\text{perm}(\mathbf{y}),\mathbf{\hat{y}}\right)$
(1)

with $\mathbf{\hat{y}}=f(\mathbf{X})$ where $f$ is our segmentation model
whose architecture is described later in the paper. In practice, for
efficiency, we first compute the $K_{\text{max}}\times K_{\text{max}}$ binary
cross entropy losses between all pairs of $\mathbf{y}$ and $\mathbf{\hat{y}}$
dimensions, and rely on the Hungarian algorithm to find the permutation that
minimizes the overall binary cross entropy loss.

###  2.2 On-the-fly data augmentation

For training, 5s audio chunks (and their reference segmentation) are cropped
randomly from the training set. To increase variability even more, we rely on
on-the-fly random data augmentation. The first type of augmentation is
additive background noise with random signal-to-noise ratio. Inspired by our
previous work on overlapped speech detection [4], the second type of
augmentation consists in artificially increasing the amount of overlapping
speech. To do that, we sum two random 5s audio chunks with random signal-to-
signal ratio (and merge their reference segmentation accordingly). Resulting
chunks whose number of speakers is higher than $K_{\text{max}}$ are not used
for training.

###  2.3 Segmentation

Once trained, the model can be used for segmentation purposes or any sub-tasks
by a simple post-processing of its output speaker activations:

  * •

for segmentation or speaker change detection, a single $\theta=0.5$
binarization threshold already gives decent results, but one can get even
better performance by using a slightly more advanced post-processing borrowed
from [14] and summarized in Figure 2.

  * •

for voice activity detection, we start by computing the maximum activation
over the $K_{\text{max}}$ speakers:

$\hat{y}_{t}^{\text{VAD}}=\max_{k}\hat{y}_{t}^{k}$ (2)

and, then only, apply the aforementioned post-processing on resulting mono-
dimensional $\mathbf{\hat{y}^{\text{VAD}}}$.

  * •

for overlapped speech detection, since at least two speakers need to be active
simultaneously to indicate overlapping speech, we compute the second highest
(denoted $\operatorname*{max_{2nd}}$) activation:

$\hat{y}_{t}^{\text{OSD}}=\operatorname*{max_{2nd}}_{k}\hat{y}_{t}^{k}$ (3)

and post-process the resulting mono-dimensional
$\mathbf{\hat{y}^{\text{OSD}}}$ with the same approach.

Figure 2: To obtain the final binary segmentation, speaker activations are
post-processed with $\theta_{\text{on}}/\theta_{\text{off}}$ hysteresis
thresholding, then filling gaps shorter than $\delta_{\text{off}}$ (light
green region in right example) and finally removing active regions shorter
than $\delta_{\text{on}}$ (does not happen in these examples). Table 1: Voice
activity detection // FA = false alarm rate / Miss. = missed detection rate

VAD AMI [8, 15] DIHARD 3 [9] VoxConverse [11] FA Miss. FA+Miss. FA Miss.
FA+Miss. FA Miss. FA+Miss. silero_vad 9.4 1.7 11.0 17.0 4.0 21.0 3.0 1.1 4.2
dihard3 [9] NA NA NA 4.0 4.2 8.2 NA NA NA Landini et al. [12] NA NA NA NA NA
NA 1.8 1.1 3.0 pyannote 1.1 [16] 6.5 1.7 8.2 4.1 3.8 7.9 4.5 0.3 4.8 Ours –
pyannote 2.0 4.0 3.2 7.2 4.6 3.6 8.2 1.9 0.8 2.7

###  2.4 Overlap-aware resegmentation

Figure 3: Effect of the proposed overlap-aware resegmentation approach (third
row) on the VBx diarization baseline (second row). We highlight three regions
where the heuristic performs better ($t\approx 100s$), same ($t\approx 120s$),
or worse ($t\approx 115$s) than the proposed approach (source: file
DH_EVAL_0035.flac in DIHARD3 dataset).

While a growing number of diarization approaches do try and take overlapped
speech into account [7], the most dependable ones (like the VBx approach [15]
used in Figure 3) still assume internally that at most one speaker is active
at any time. There is therefore a need for a post-processing step that assigns
multiple speaker labels to overlapped speech regions [4, 17].

Given an existing speaker diarization output (with $K$ speakers) encoded into
a sequence of $K$-dimensional binary frames $y_{t}^{\text{DIA}}$, we propose
to use the segmentation model as a local, overlap-aware, resegmentation
module. The segmentation model is applied on a 5s-long window sliding over the
whole file. At each step, we find the permutation of the speaker activations
$\mathbf{\hat{y}}$ that minimizes the binary cross entropy loss with respect
to $\mathbf{y}^{\text{DIA}}$. Permutated sliding speaker activations are then
aggregated over time and post-processed with the threshold-based approach
introduced in Section 2.3.

##  3 Experiments

Datasets and partitions.  We ran experiments and report results on three
speaker diarization datasets, covering a wide range of domains:

DIHARD3 corpus [9, 10] does not provide a training set. Therefore, we split
its development set into two parts: 192 files used as training set, and the
remaining 62 files used as a smaller development set. The latter is simply
referred to as development sets in the rest of the paper. When defining this
split (shared at huggingface.co/pyannote/segmentation), we made sure that the
11 domains were equally distributed between both subsets. The evaluation set
is kept unchanged.

VoxConverse does not provide a training set either [11]. Therefore, we also
split its development set into two parts: first 144 files (abjxc to qouur, in
alphabetical order) constitute the training set, leaving the remaining 72
files (qppll to zyffh) for the actual development set.

AMI provides an official {training, development, evaluation} partition of the
Mix-Headset audio files [8]. While we kept the development and evaluation sets
unchanged, we only used the first 10 minutes of each file of the training set,
to end up with an actual training set similar in size (22 hours) to that of
the DIHARD3 (25 hours) and VoxConverse (15 hours) training sets.

Experimental protocols.  We trained a unique segmentation model using the
composite training set (62 hours) made of the concatenation of all three
training sets. The composite development set (24 hours) served as validation
and was used to decrease the learning rate on plateau and eventually choose
the best model checkpoint. At the end of this process, only one segmentation
model is available (not one model per dataset) and used for all experiments.

However, detection thresholds ($\theta_{\text{on}}$, $\theta_{\text{off}}$,
$\delta_{\text{on}}$, and $\delta_{\text{off}}$) were tuned specifically for
each dataset using their own development set because the manual annotation
guides differ from one dataset to another, especially regarding
$\delta_{\text{off}}$ which controls whether to bridge small intra-speaker
pauses. For the same reasons, detection thresholds were optimized specifically
for each task addressed in the paper:

  * •

voice activity detection thresholds are chosen to minimize the detection error
rate (i.e. the sum of the false alarm and missed detection rates), with no
forgiveness collar around speech turn boundaries;

  * •

overlapped speech detection thresholds are chosen to maximize the detection
$F_{1}$-score, with no forgiveness collar either;

  * •

for resegmentation, detection thresholds are chosen to minimize the
diarization error rate, without forgiveness collar but with overlapped speech
regions. This is consistent with DIHARD3 evaluation plan [10] and AMI Full
evaluation setup [15], but not with VoxConverse challenge rules that uses a
250ms collar [11].

All metrics were computed using pyannote.metrics [18] open source Python
library.

Implementation details.  Our segmentation model ingests 5s audio chunks with a
sampling rate of 16kHz (i.e. sequences of 80000 samples). The input sequence
is passed to SincNet convolutional layers using the original configuration
[19] – except for the stride of the very first layer which is set to 10 (so
that SincNet frames are extracted every 16ms). Four bidirectional Long Short-
Term Memory (LSTM) recurrent layers (each with 128 units in both forward
backward directions, and 50% dropout for the first three layers) are stacked
on top of two additional fully connected layers (each with 128 units and leaky
ReLU activation) which also operate at frame-level. A final fully connected
classification layer with sigmoid activation outputs
$K_{\text{max}}$-dimensional speaker activations between 0 and 1 every 16ms.
Overall, our model contains 1.5 million trainable parameters – most of which
(1.4 million) comes from the recurrent layers.

As introduced in Section 2.2, 50% of the training samples are made out of the
weighted sum of two chunks, with a signal-to-signal ratio sampled uniformly
between 0 and 10dB. We also use additive background noise from the MUSAN
dataset [20] with a signal-to-noise ratio sampled uniformly between 5 and
15dB.

We train the model with Adam optimizer with default PyTorch parameters and
mini-batches of size 128. Learning rate is initialized at $10^{-3}$ and
reduced by a factor of 2 every time its performance on the development set
reaches a plateau. It took around 3 days using 4 V100 GPUs to reach peak
performance. While we do share the pretrained model at
huggingface.co/pyannote/segmentation for reproducing the results, the whole
training process is also reproducible as everything has been integrated into
version 2.0 of pyannote.audio open-source library [16].

##  4 Results and discussions

Table 2: Overlapped speech detection // $F_{1}$ = $F_{1}$-score

OSD AMI [8, 15] DIHARD 3 [9] VoxConverse [11] Precision Recall $F_{1}$
Precision Recall $F_{1}$ Precision Recall $F_{1}$ Kunesova et al. [3] 71.5
46.1 56.0 NA NA NA NA NA Landini et al. [12] NA NA NA NA NA NA 73.0 28.2 40.7
pyannote 1.1 [16, 4] 63.2 87.9 73.5 53.2 54.8 54.0 38.7 82.3 52.6 Ours –
pyannote 2.0 77.1 70.9 73.9 54.3 61.1 57.5 69.1 77.6 73.1

Table 3: Resegmentation // FA = false alarm / Miss. = missed detection / Conf.
= speaker confusion / DER = diarization error rate

Baseline Overlap-aware AMI [8, 15] DIHARD 3 [9] VoxConverse [11]
resegmentation FA Miss. Conf. DER FA Miss. Conf. DER FA Miss. Conf. DER
pyannote 1.1 [16] _ 5.0 16.2 8.5 29.7 3.4 13.2 12.6 29.2 2.0 10.1 9.5 21.5
Heuristic [13] w/ our OSD 7.3 7.8 10.9 26.0 6.6 9.0 12.7 28.3 3.1 7.3 10.1
20.5 Ours – pyannote 2.0 4.0 13.0 9.1 26.1 5.1 9.8 10.3 25.2 2.4 3.1 9.8 15.4
dihard3 [9] _ NA NA NA NA 3.6 13.3 8.4 25.4 NA NA NA NA Heuristic [13] w/ our
OSD NA NA NA NA 7.0 8.8 8.8 24.6 NA NA NA NA Ours – pyannote 2.0 NA NA NA NA
4.6 10.2 7.5 22.2 NA NA NA NA VBx [15] w/ our VAD _ 3.4 17.3 3.9 24.7 4.2 12.7
6.0 23.0 1.8 5.5 1.2 8.5 Heuristic [13] w/ our OSD 5.9 8.7 6.3 21.0 7.9 8.1
6.2 22.2 3.0 2.4 1.8 7.2 Ours – pyannote 2.0 4.5 11.0 4.8 20.3 4.7 9.7 4.7
19.2 2.7 2.6 1.8 7.1 Oracle Ours – pyannote 2.0 4.7 10.0 1.4 16.1 4.6 9.8 1.8
16.2 2.6 2.5 0.6 5.7

Voice activity detection.  Table 1 compares the performance of the proposed
voice activity detection approach with the official dihard3 baseline [9],
Landini’s submission to VoxConverse challenge [12], and pyannote 1.1 VAD
models [16]. The main conclusion is that, despite it being trained for
segmentation, our model can still be competitive with (if not better than)
models trained specifically for voice activity detection. Note, however, that
one should not draw hasty conclusions regarding the performance of silero_vad
model [21] as it is an off-the-shelf model which was not trained specifically
for these datasets.

Overlapped speech detection.  Finding good and reproducible baselines for the
overlapped speech detection task proved to be a difficult task. We thank
Kunesova et al. [3] and Landini et al. [12] for sharing the output of their
detection pipelines. Results are reported in Table 2 that shows that, like for
voice activity detection, our segmentation model can be used successfully for
overlapped speech detection, even though it was not initially trained for this
particular task. It outperforms pyannote 1.1 overlapped speech detection which
we believe was the previous state of the art [4].

Overlap-aware resegmentation.  While our segmentation model was found to be
useful for both voice activity detection and overlapped speech detection,
post-processing the output of existing speaker diarization pipelines is where
it really shines. Table 3 summarizes the resegmentation experiments performed
on top of three of them, ranked from worst to best baseline performance:
pyannote 1.1 pretrained pipelines [16], dihard3 official baseline [9], and
BUT’s VBx approach [15]. The (admittedly wrong) criterion used for selecting
those baselines was their ease of use and reproducibility. Because results
reported in [15] for VBx baseline rely on oracle voice activity detection and
the shared code base does not provide an official voice activity detection
implementation, we used our own (marked as Ours in Table 1) and applied VBx on
top of it. Our proposed resegmentation approach consistently improves the
output of all baselines on all datasets. Relative diarization error rate
improvement over the best baseline (VBx) reaches 18% on AMI, 17% on DIHARD,
and 16% on VoxConverse.

For comparison purposes, we also implemented a heuristic that consists in
assigning detected overlapped speech regions to the two nearest speakers in
time [13]. Despite its simplicity, this heuristic happens to be a strong
baseline, very difficult to beat in practice [12]. Yet, our proposed
resegmentation approach outperforms it for all but one experimental condition
(for which the heuristic is better only by a small margin). A closer look at
the speaker confusion error rates shows that our approach is significantly
better at identifying overlapping speakers. This is confirmed by the low
speaker confusion error rates obtained when we apply it on top of an oracle
diarization (with $\mathbf{y}^{\text{DIA}}=\mathbf{y}$): only 1.4%, 1.8%, and
0.6% of speech are re-assigned incorrectly on AMI, DIHARD and VoxConverse
respectively. Figure 3 provides a qualitative sneak peak at their respective
behavior on a short 20 seconds excerpt. In particular, it appears that the two
(heuristic and proposed) approaches do behave differently and could complement
each other.

##  5 Conclusions

The overall best pipeline reported in this paper is the combination of our
voice activity detection, off-the-shelf VBx clustering, and our overlap-aware
resegmentation approach, reaching $\text{DER}=20.3\%$ on AMI Mix-Headset using
the full evaluation setup introduced in [15], $\text{DER}=19.2\%$ on DIHARD 3
evaluation set (full condition, 2.5% behind the winning submission), and
$\text{DER}=7.1\%$ (or $\text{DER}=3.4\%$ with a 250ms forgiveness collar) on
VoxConverse development set.

Even with a forgiveness collar, missed detection and false alarms are the main
source of errors (twice as high as speaker confusion) for all three datasets –
indicating that, despite progress, overlapped speech detection remains an
unsolved (and sometimes ill-defined) problem. A low hanging fruit could be to
replace recurrent layers with Transformers as proposed in [6]. One could also
add a second term in the loss function defined in Equation 1, focusing on
overlapped speech detection and inspired by Equation 3. The lazy authors are
even willing to help the interested reader integrate these ideas into the
pyannote.audio code base.

## References

  * [1] R. Yin, H. Bredin, and C. Barras, “Speaker Change Detection in Broadcast TV Using Bidirectional Long Short-Term Memory Networks,” in _Proc. Interspeech 2017_ , 2017.
  * [2] D. Charlet, C. Barras, and J. Liénard, “Impact of overlapping speech detection on speaker diarization for broadcast news and debates,” in _2013 IEEE International Conference on Acoustics, Speech and Signal Processing_ , May 2013, pp. 7707–7711.
  * [3] M. Kunešová, M. Hrúz, Z. Zajíc, and V. Radová, “Detection of overlapping speech for the purposes of speaker diarization,” in _Speech and Computer_ , A. A. Salah, A. Karpov, and R. Potapova, Eds. Cham: Springer International Publishing, 2019, pp. 247–257.
  * [4] L. Bullock, H. Bredin, and L. P. Garcia-Perera, “Overlap-aware diarization: Resegmentation using neural end-to-end overlapped speech detection,” in _Proc. ICASSP 2020_ , 2020.
  * [5] Y. Fujita, N. Kanda, S. Horiguchi, K. Nagamatsu, and S. Watanabe, “End-to-End Neural Speaker Diarization with Permutation-free Objectives,” in _Interspeech_ , 2019, pp. 4300–4304.
  * [6] Y. Fujita, N. Kanda, S. Horiguchi, Y. Xue, K. Nagamatsu, and S. Watanabe, “End-to-end neural speaker diarization with self-attention,” in _2019 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU)_ , 2019, pp. 296–303.
  * [7] Y. Takashima, Y. Fujita, S. Watanabe, S. Horiguchi, P. García, and K. Nagamatsu, “End-to-end speaker diarization conditioned on speech activity and overlap detection,” in _2021 IEEE Spoken Language Technology Workshop (SLT)_ , 2021, pp. 849–856.
  * [8] J. Carletta, “Unleashing the killer corpus: experiences in creating the multi-everything AMI Meeting Corpus,” _Language Resources and Evaluation_ , vol. 41, no. 2, 2007.
  * [9] N. Ryant, P. Singh, V. Krishnamohan, R. Varma, K. Church, C. Cieri, J. Du, S. Ganapathy, and M. Liberman, “The Third DIHARD Diarization Challenge,” _arXiv preprint arXiv:2012.01477_ , 2020.
  * [10] N. Ryant, K. Church, C. Cieri, J. Du, S. Ganapathy, and M. Liberman, “Third DIHARD Challenge Evaluation Plan,” _arXiv preprint arXiv:2006.05815_ , 2020.
  * [11] J. S. Chung, J. Huh, A. Nagrani, T. Afouras, and A. Zisserman, “Spot the Conversation: Speaker Diarisation in the Wild,” in _Proc. Interspeech 2020_ , 2020, pp. 299–303. [Online]. Available: http://dx.doi.org/10.21437/Interspeech.2020-2337
  * [12] F. Landini, O. Glembek, P. Matějka, J. Rohdin, L. Burget, M. Diez, and A. Silnova, “Analysis of the BUT Diarization System for VoxConverse Challenge,” in _2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2021.
  * [13] S. Otterson and M. Ostendorf, “Efficient use of overlap information in speaker diarization,” in _2007 IEEE Workshop on Automatic Speech Recognition & Understanding (ASRU)_. IEEE, 2007, pp. 683–686.
  * [14] G. Gelly and J.-L. Gauvain, “Optimization of RNN-Based Speech Activity Detection,” _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , vol. 26, no. 3, pp. 646–656, March 2018.
  * [15] F. Landini, J. Profant, M. Diez, and L. Burget, “Bayesian hmm clustering of x-vector sequences (vbx) in speaker diarization: theory, implementation and analysis on standard tasks,” 2020.
  * [16] H. Bredin, R. Yin, J. M. Coria, G. Gelly, P. Korshunov, M. Lavechin, D. Fustes, H. Titeux, W. Bouaziz, and M.-P. Gill, “pyannote.audio: neural building blocks for speaker diarization,” in _Proc. ICASSP 2020_ , 2020.
  * [17] S. Horiguchi, P. Garcia, Y. Fujita, S. Watanabe, and K. Nagamatsu, “End-to-end speaker diarization as post-processing,” 2020.
  * [18] H. Bredin, “pyannote.metrics: a toolkit for reproducible evaluation, diagnostic, and error analysis of speaker diarization systems,” in _Proc. Interspeech 2017_ , Stockholm, Sweden, August 2017. [Online]. Available: http://pyannote.github.io/pyannote-metrics
  * [19] M. Ravanelli and Y. Bengio, “Speaker recognition from raw waveform with sincnet,” in _Proc. SLT 2018_ , 2018.
  * [20] D. Snyder, G. Chen, and D. Povey, “MUSAN: A Music, Speech, and Noise Corpus,” 2015.
  * [21] Silero Team, “Silero VAD: pre-trained enterprise-grade Voice Activity Detector (VAD), Number Detector and Language Classifier,” https://github.com/snakers4/silero-vad, 2021.


# Balancing Approach for Causal Inference at Scale

Sicheng Lin  Snap Inc.Santa MonicaUSA90405<EMAIL_ADDRESS>,  Meng Xu  Snap
Inc.Santa MonicaUSA90405<EMAIL_ADDRESS>,  Xi Zhang  Snap Inc.Santa
MonicaUSA90405<EMAIL_ADDRESS>,  Shih-Kang Chao  University of
MissouriColumbiaMOUSA<EMAIL_ADDRESS>,  Ying-Kai Huang  Realtor.comSanta
ClaraCAUSA<EMAIL_ADDRESS>and  Xiaolin Shi  Snap Inc.Santa
MonicaUSA90405<EMAIL_ADDRESS>
(2023)

###### Abstract.

With the modern software and online platforms to collect massive amount of
data, there is an increasing demand of applying causal inference methods at
large scale when randomized experimentation is not viable. Weighting methods
that directly incorporate covariate balancing have recently gained popularity
for estimating causal effects in observational studies. These methods reduce
the manual efforts required by researchers to iterate between propensity score
modeling and balance checking until a satisfied covariate balance result.
However, conventional solvers for determining weights lack the scalability to
apply such methods on large scale datasets in companies like Snap Inc. To
address the limitations and improve computational efficiency, in this paper we
present scalable algorithms, DistEB and DistMS, for two balancing approaches:
entropy balancing (Hainmueller, 2012) and MicroSynth (Robbins et al., 2017).
The solvers have linear time complexity and can be conveniently implemented in
distributed computing frameworks such as Spark, Hive, etc. We study the
properties of balancing approaches at different scales up to 1 million treated
units and 487 covariates. We find that with larger sample size, both bias and
variance in the causal effect estimation are significantly reduced. The
results emphasize the importance of applying balancing approaches on large
scale datasets. We combine the balancing approach with a synthetic control
framework and deploy an end-to-end system for causal impact estimation at Snap
Inc.

causal inference; observational study; big data; covariate balancing;
quadratic programming

††journalyear: 2023††copyright: acmlicensed††conference: Proceedings of the
29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining; August
6–10, 2023; Long Beach, CA, USA††booktitle: Proceedings of the 29th ACM SIGKDD
Conference on Knowledge Discovery and Data Mining (KDD ’23), August 6–10,
2023, Long Beach, CA, USA††price: 15.00††doi: 10.1145/3580305.3599778††isbn:
979-8-4007-0103-0/23/08††ccs: Computing methodologies MapReduce
algorithms††ccs: General and reference Performance††ccs: General and reference
Estimation

##  1\. Introduction

Within the flourishing era of big data, we are dealing with an explosion of
available information and high dimensional data sets. Companies like Snap Inc.
conduct many quasi-experimental studies with millions of observations, aiming
to demonstrate the causal relationship between an intervention and tens of
critical business metrics. These causal studies boost the data-driven fashion
of decision making processes and have the power to pivot future strategic
business decisions. They are crucial to the company yet cannot be AB tested
due to the reason that we are unable to intervene users in these cases.

In observational studies, it is common that some of the covariates are
imbalanced between treated group and control group because the treatment
conditions are not randomly assigned to units. Conducting a causal analysis
becomes difficult when the covariates are associated with both the treatment
assignment and the outcome, as the imbalances in these covariates may
introduce selection bias to the treatment effect estimation. In order to
adjust for covariate imbalance, the pioneering work of propensity score
(Rosenbaum and Rubin, 1983) demonstrates a key point that treatment assignment
is independent of observed covariates conditional on the propensity score.
Based on this property, many statistical methods has been proposed using
propensity score, including propensity score matching (Rosenbaum and Rubin,
1983, 1985), propensity score stratification (Rosenbaum and Rubin, 1984),
inverse probability weighting (Robins et al., 1994; Hirano and Imbens, 2001),
etc. Among these methods, weighting is the one with the most attractive
statistical property. First, if propensity score is correctly modeled, the
weighting estimator is consistent (Lunceford and Davidian, 2004). Second,
nonparametric propensity score weighting estimator is efficient (Hirano et
al., 2003). As a comparison, matching and stratification estimators do not
have these properties (Abadie and Imbens, 2006; Lunceford and Davidian, 2004).
Furthermore, weighting methods do not require explicit outcome modeling, the
weights are obtained solely from user covariates in the experimental design
phase. One set of weights can be used to estimate causal effects for multiple
outcomes. In this regard, the weighting methods are similar to A/B testing
(Rubin, 2008) as they are considered more of experimental design methods than
experimental analysis methods.

Despite the solid theoretical foundation of the propensity score weighting
method, in practical usage the performance of this approach is sensitive to
model misspecfication. When applying propensity score weighting methods,
applied researchers tend to alternately tune propensity score model and
conduct balance test until a satisfied balacing result (Imai et al., 2008).
Entropy Balancing (hereafter EB) (Hainmueller, 2012), MicroSynth (Robbins et
al., 2017) and several other methods (Zubizarreta, 2015; Imai and Ratkovic,
2014; Zhao, 2019) have been proposed recently to overcome the problem by
directly balancing covariates instead of modeling propensity scores. In this
paper we refer to this class of methods as balancing approach. In general,
balancing approach works better than propensity score modeling approach by
achieving better balance among covariates. It is also shown that EB has doubly
robust property under linear model specification (Zhao and Percival, 2017).

We rely on balancing approaches for solving causal inference problems because
a set of weights can be used repeatedly for multiple business metrics and
balancing approaches normally have smaller bias than propensity score modeling
approaches. However we find that the current balancing approaches lack the
computational efficiency to handle such big data. The number of units and
dimensions in our applications can be as large as 100 millions and more than
1000 dimensions, respectively, and the overall data size can reach several
tera-bytes. The majority of the existing causal inference software packages
are designed for problems where the data set can be fit into one single
machine. They are implemented mostly in data-science friendly languages such
as Python or R, rarely to be found using distributed framework that suitable
for large scale data. A Stata package Ebalance (Hainmueller and Xu, 2013) and
a R package ebal (Hainmueller, 2022) are the first software to find the
entropy balancing weights. The microsynth (Robbins and Davenport, 2021)
package, implemented in R, finds MicroSynth weights by solving a quadratic
programming problem. Later on empirical_calibration (hereafter EC) (Wang et
al., 2019), a Python library developed by Google that can solve both entropy
objective and quadratic objective, improves the computation speed of ebal by
22 times. Aside from balancing approaches, other modern causal inference
software such as EconML (Battocchi et al., 2019) and DoWhy (Sharma and
Kiciman, 2020) by Microsoft, CausalML (Chen et al., 2020) by Uber,
CausalImpact (Brodersen et al., 2014) by Google, create a rich array of causal
tooling yet are limited by the memory constraint from computer machine. There
is a lack of mature software dedicated for large scale causal inference tasks
(Wong, 2020).

In order to scale up causal inference methods, Netflix invested in data
compression strategies (Wong et al., 2021) by extracting the summary
statistics from original datasets. They find this strategy is able to
losslessly estimate parameters from a large class of linear models (Wong et
al., 2021). Several other works (e.g. (Ouattara et al., 2021; Bluhm and
Cutura, 2022)) attempt to use parallel computing to scale up linear models
including OLS, logistic regressions etc. There is an absence of studies to
scale up more advanced causal methods like balancing approaches. We design
algorithms that directly utilize distributed computing frameworks to solve the
optimization problems in balancing approaches. To the best of our knowledge,
this is the first work that scales balancing approaches up to large datasets.

The contribution of this paper is twofold. First, we present two distributed
computing algorithms, which we name as DistEB and DistMS, to solve two well-
known balancing approaches: Entropy Balancing and MicroSynth. We exploit the
dual function of the optimization problem in both approaches. We use MapReduce
framework to calculate the gradient of Lagrange multiplier and apply iterative
gradient based solvers. Our algorithm scales linearly with the number of
observations as well as the number of covariates. It is feasible to estimate
causal effects up to 100 millions units and more than 1000 dimensions. Second,
we apply our methods and find that larger scale datasets significantly reduce
both the bias and the variance in causal effect estimations, highlighting the
importance of scalability of causal inference methods.

The rest of this paper is structured as follows. In Section 2 we introduce the
causal inference assumptions and the optimization problems in balancing
approach. In Section 3 we describe our algorithms DistEB and DistMS in detail.
Section 4 evaluates the performance of EB and MicroSynth as well as the
scalabity of our algorithms. We discuss the applications of our algorithms at
Snap Inc. in Section 5 and conclude in Section 6.

##  2\. Preliminaries

###  2.1. Causal Inference with Observational Data: Setup and Assumptions

Let us consider the following causal inference problem in observational
studies. In a random sample of size $N$ drawn from a population, each unit $i$
receives a binary treatment assignment $T_{i}$. $T_{i}=1$ if unit $i$ is
assigned in the treated group and $T_{i}=0$ if unit $i$ is in the control
group. The number of units in the control group and the treated group are
$N_{0}$ and $N_{1}$, respectively. We observe an outcome $Y_{i}$ and a set of
covariates $X_{i}=[X_{i1},X_{i2},\cdots,X_{id}]^{\top}$ for each unit $i$,
where $d$ is the dimension of covariates. Assuming the Stable Unit Treatment
Value Assumption (SUTVA) (Rubin, 1980), the potential outcome
$\\{Y_{i}^{(0)},Y_{i}^{(1)}\\}$ does not depend on the treatment assignment of
other units, i.e., no network effects. Then the observed outcome can be
expressed as $Y_{i}=T_{i}Y_{i}^{(1)}+(1-T_{i})Y_{i}^{(0)}$.

The treatment effect of each unit is defined as
$\tau_{i}=Y_{i}^{(1)}-Y_{i}^{(0)}$. In this paper we focus on Population
Average Treatment Effect on the Treated (PATT) estimation,

(1)
$\tau_{\textrm{PATT}}=\mathbb{E}[\tau_{i}|T_{i}=1]=\mathbb{E}[Y_{i}^{(1)}|T_{i}=1]-\mathbb{E}[Y_{i}^{(0)}|T_{i}=1].$

The first expectation can be simply estimated as the sample mean of outcome in
the treated group. The second term describes unobserved counterfactual and
needs to be estimated using data from the control group. In a randomized
controlled trial, the treatment groups are randomly sampled from a population,
therefore we can use the sample mean of outcome in the control group to
estimate the second term because
$\mathbb{E}[Y_{i}^{(0)}|T_{i}=1]=\mathbb{E}[Y_{i}^{(0)}|T_{i}=0]$. However in
observational studies, where the treatment assignment is not random, selection
into treatment and the outcome may both relate to confounding factors, causing
the two expectations unequal. In order to correct the selection bias caused by
confounders, we adopt the assumptions of strong ignorability and overlap from
(Rosenbaum and Rubin, 1983).

Assumption 1 (strong ignorability): $(Y^{(0)},Y^{(1)})\perp T|X$. It implies
that there are no unobserved confounders.

Assumption 2 (overlap): $0<P(T=1|X)<1$. It states that every unit has positive
possibility of receiving each treatment condition.

###  2.2. Weighting Methods: Modeling Approach VS Balancing Approach

Originated from inverse probability weighting (IPW), weighting methods attempt
to find weight for each unit such that the imbalance among observed covariates
are adjusted. For the PATT estimation, only units in the control group need to
be assigned with weights. The weighting estimator is calculated as:

(2)
$\widehat{\tau_{\textrm{PATT}}}=\sum_{T_{i}=1}\frac{1}{N_{1}}Y_{i}-\sum_{T_{i}=0}w_{i}Y_{i},$

where $w_{i}$ is the weight for unit $i$ and the sum of weights satisfy
$\sum_{T_{i}=0}w_{i}=1$, $N_{1}$ is the number of treated units.

There are two general approaches to calculate the weights. The modeling
approach estimates the weights by modeling the probabilities of receiving
treatment. The balancing approach directly minimizes covariate imbalance and
the dispersion of weights.

####  2.2.1. Modeling Approach

The modeling approach is built upon the propensity score (Rosenbaum and Rubin,
1983), defined as the probability of treatment assignment conditional on
observed covariates. In observational studies, this quantity is unknown and
need to be estimated by fitting a selection model on the observed data.

There are many ways to fit a propensity score model. The most popular one is
logistic regression. Propensity score model can also be based on non-linear or
non-parametric methods including regularized logistic regression, random
forest, boosted regression, support vector machines, neural networks,
multivariate adaptive regression splines, etc.

To estimate the PATT, the weight in IPW method is

(3)
$w_{i}^{\textrm{IPW}}=\frac{\hat{e}(X_{i})(1-\hat{e}(X_{i}))^{-1}}{\sum_{T_{i}=0}\hat{e}(X_{i})(1-\hat{e}(X_{i}))^{-1}},$

where $\hat{e}(X_{i})$ is the estimated probability of receiving treatment.
When the propensity scores are estimated consistently, the IPW estimator is
also consistent.

One way to improve IPW estimator is by combining the propensity score model
with an outcome model to achieve doubly robustness (DR). This doubly robust
procedure produces consistent estimates if either of the two models are
correctly specified (Robins and Rotnitzky, 1995). A DR estimator can be
constructed in numerous ways (Kang and Schafer, 2007). A modified form of
regression estimation with residual bias correction employed by (Zhao and
Percival, 2017) is

(4)
$\widehat{\tau_{\textrm{PATT}}^{\textrm{IPW,DR}}}=\sum_{T_{i}=1}\frac{1}{N_{1}}(Y_{i}-\hat{g_{0}}(X_{i}))-\sum_{T_{i}=0}w_{i}^{\textrm{IPW}}(Y_{i}-\hat{g_{0}}(X_{i})),$

where $\hat{g_{0}}$ is an outcome model based on the data in the control
group.

There are several limitations of the modeling approach. First, when the
propensity score model is correctly specified, the IPW weights only balance
covariates in expectation but not in finite sample. In any particular finite
sample, it could be difficult to balance covariates. Second, balance check is
necessary for the modeling approach. If covariates remain imbalanced between
the treatment groups after IPW adjustment, we need to modify the propensity
score model to achieve satisfactory balance. This adjustment procedure
involves in model selection and hyperparameter tuning. With big data, it could
be very expensive to look for the model to achieve the best balance. Because
of the first limitation, the model achieves the best balance in a finite
sample is not necessarily the correctly specified model or the model with the
smallest bias (see discussion in Section 4.3). It means the expensive model
selection and tuning process does not guarantee satisfactory performance. The
doubly robust IPW estimator in Equation 4 has much better performance than the
IPW estimator in our simulated data. However, when we have a large number of
outcome variables, we need to model them separately. The doubly robust IPW
method could be expensive in big data with a large number of outcome
variables.

####  2.2.2. Balancing Approach

The balancing approach has been proposed to overcome the limitations of the
modeling approach. This approach directly balance the covariates without
explicitly modeling the probability of treatment. Entropy Balancing (EB)
(Hainmueller, 2012) is a groundbreaking method that minimize the Kullback-
Leibler divergence between estimated weights and base weight. Usually the base
weights are set to uniform weights and can be omitted from the objective
function. The EB weights are calculated by solving the following optimization
problem

(5) $\displaystyle\min$ $\displaystyle\sum_{T_{i}=0}w_{i}\log(w_{i})$ s.t.
$\displaystyle\sum_{T_{i}=0}w_{i}c_{j}(X_{i})=\frac{1}{N_{1}}\sum_{T_{i}=1}c_{j}(X_{i})\
\forall j,$ $\displaystyle\sum_{T_{i}=0}w_{i}=1,$ $\displaystyle w_{i}\geq 0\
\forall i,$

where functions $c_{j}$ are moment functions of the covariates. Substituting
the solution of the above optimization problem into Equation 2, the EB
estimator of PATT is

(6)
$\widehat{\tau_{\textrm{PATT}}^{\textrm{EB}}}=\sum_{T_{i}=1}\frac{1}{N_{1}}Y_{i}-\sum_{T_{i}=0}w_{i}^{\textrm{EB}}Y_{i}.$

By enforcing covariate balancing constraint, EB implicitly fits a linear
regression model. The EB estimator is proved to be doubly robust and achieves
the efficiency bound (Zhao and Percival, 2017). This indicates that if either
the true propensity score model or outcome model are linear, the PATT is
consistently estimated. Simulation results in Section 4 shows that EB has much
better performance in balance and bias than IPW and similar performance to
doubly robust IPW.

There are several extensions based on EB. Instead of the K-L divergence, the
objective function of the minimization problem can also be other forms,
including empirical balancing calibration weights with $D(w_{i},1)$ (Chan et
al., 2016),111$D(x,x_{0})$ is a distance measure for a fixed
$x_{0}\in\mathbb{R}$ that is continuously differentiable in
$x_{0}\in\mathbb{R}$, nonnegative and strictly convex in $x$. the stable
balancing weights which limits the possibility of extreme weights with
$(w_{i}-1/r)^{2}$ (Zubizarreta, 2015), and a special case of stable balancing
weights, MicroSynth with $\frac{1}{2}(w_{i}-1)^{2}$ (Robbins et al., 2017). If
the constraints in the optimization problem 5 hold, these extensions are also
doubly robust (Zhao and Percival, 2017).

Another popular method of balancing approach is covariate balancing propensity
score (CBPS) (Imai and Ratkovic, 2014). In contrast to EB and its extensions,
CBPS explicitly fits the propensity score model with balancing weights. This
method is a combination of modeling approach and balancing approach. However,
our simulation shows CBPS may not achieve satisfactory balance. It leaves the
estimated PATT susceptible to bias.

It is noteworthy to mention that the balancing approach does not limit to
estimate PATT. Some recent studies extend this framework to estimate the
Average Treatment Effect (ATE) (Josey et al., 2021; Chen et al., 2022).

More advanced methods of balancing approach emerged recently. Kernel balancing
(Hazlett, 2020) makes the estimated multivariate density of covariates
approximately the same for the treated and control groups by balancing the
kernel matrix instead of original covariates . Energy balancing (Huling and
Mak, 2020) aims to balance the distribution of covariates. However, these
methods are much more computationally expensive than EB and its extensions
therefore pose a practical limit of sample size and dimensions.

##  3\. Balancing Approach at Scale

Our main objective in this paper is to provide a scalable solution for
balancing approaches such as EB and MicroSynth. The convex optimization
problem in balancing approaches can be summarised in Equation 5 with different
forms of loss function. The original paper of EB (Hainmueller, 2012) and the
implementation (Hainmueller and Xu, 2013) suggest solving the optimization
problem by Newton method with a line search on the Lagrange dual problem.
Newton methods require calculations of the inverse of the Hessian matrix for
the parameter update direction. However the inverse matrix calculation can be
computationally expensive in high dimensional settings. Similar computational
concerns arise for MicroSynth, where the objective function is a square loss
and the optimization problem becomes a high dimensional quadratic programming
with both equality and inequality constraints. We aim to provide a solution
that scales linearly with number of units and utilize distributed computing
framework such as Spark or Hadoop.

###  3.1. Entropy Balancing

For notation simplicity, we consider the first moment function
$c_{j}(X_{i})=X_{i}$ in the following sections, the results can be applied for
higher moments as well. In order to solve the EB optimization problem, we
apply Lagrange multipliers $\bm{\xi}\in\mathbb{R}^{d}$ to obtain the primal
problem:

(7)
$\textrm{min}\bigg{\\{}\mathcal{L}^{\textrm{EB}}=\sum_{T_{i}=0}w_{i}\log(w_{i})+\bm{\xi}^{\top}(\sum_{T_{i}=0}w_{i}X_{i}-\tilde{X})+\xi_{0}(\sum_{T_{i}=0}w_{i}-1)\bigg{\\}},$

where $\tilde{X}=\frac{1}{N_{1}}\sum_{T_{i}=1}X_{i}$ is the treatment mean of
covariates, $\xi_{0}$ is the Lagrange multiplier on the normalization
constraint. Solving $\partial\mathcal{L}^{\textrm{EB}}/\partial w_{i}=0$
yields solution of $w_{i}$:

(8)
$\hat{w}_{i}=w_{i}(\hat{\bm{\xi}})=\frac{\exp(-\hat{\bm{\xi}}^{\top}X_{i})}{\sum_{T_{i}=0}\exp(-\hat{\bm{\xi}}^{\top}X_{i})}.$

Notice that this expression automatically satisfy the non-negative condition
of $w_{i}\geq 0$. Given $\hat{w}_{i}$, the Lagrange dual problem is

(9)
$\min_{\bm{\xi}\in\mathbb{R}^{d}}\bigg{\\{}\mathcal{L}(\bm{\xi})=\log\bigg{(}\sum_{T_{i}=0}\exp(-\bm{\xi}^{\top}X_{i})\bigg{)}+\bm{\xi}^{\top}\tilde{X}\bigg{\\}}.$

Considering the case where the data size is large and cannot be fitted into
one single machine, we seek for distributed methods to solve the optimization
problem in Equation 9. While implementing large scale distributed L-BFGS
solver remains an option for us, we utilize gradient descent solver for its
simplicity and computation efficiency. The gradient of the dual problem can be
expressed as

(10)
$\frac{\partial\mathcal{L}}{\partial\bm{\xi}}=\tilde{X}-\sum_{T_{i}=0}\frac{\exp(-\bm{\xi}^{\top}X_{i})}{\sum_{T_{i}=0}\exp(-\bm{\xi}^{\top}X_{i})}X_{i}=\tilde{X}-\sum_{T_{i}=0}w_{i}(\bm{\xi})X_{i}.$

At each updating step of gradient descent, we first compute the weights
$w_{i}(\bm{\xi})$ using MapReduce, the gradient can then be calculated by a
second MapReduce step. In practice, momentum (Qian, 1999) is used to
accelerate the convergence rate. Once $\bm{\xi}$ is solved from the dual
problem, the final weights in Equation 8 can be easily calculated using
MapReduce. We present DistEB in Algorithm 1.

Algorithm 1 Distributed Entropy Balancing (DistEB)

1:Randomly initialize
$\bm{\xi}^{(0)}\leftarrow\mathcal{N}(\mathbf{0},\bm{I}_{d})$, initialize
learning rate $\alpha\leftarrow 0.01,\beta\leftarrow 0.9$.

2:for $k=1,\ldots,K$ do $\triangleright$ $K$: the number of iterations

3: Distributed compute weights $w_{i}(\bm{\xi}^{(k-1)})$ defined in Equation
8. $\triangleright$ MapReduce

4:
$\partial\mathcal{L}/\partial\bm{\xi}^{(k-1)}=\tilde{X}-\sum_{T_{i}=0}w_{i}\big{(}\bm{\xi}^{(k-1)}\big{)}X_{i}$
$\triangleright$ MapReduce

5:
$\bm{v}^{(k)}\leftarrow-\alpha\partial\mathcal{L}/\partial\bm{\xi}^{(k-1)}+\beta*\bm{v}^{(k-1)}$
$\triangleright$ Momentum step

6: $\bm{\xi}^{(k)}\leftarrow\bm{\xi}^{(k-1)}+\bm{v}^{(k)}$ $\triangleright$
Gradient descent

7:end for

8:Compute final weights $w_{i}^{\textrm{EB}}=w_{i}(\bm{\xi}^{(k)})$ for all
control users. $\triangleright$ MapReduce

We iterate through the whole data set twice at each updating step, the overall
time complexity is $O(NdK)$, where $N$ is number of control units, $d$ is the
dimensions of balancing covariates, and $K$ is the number of iterations.
DistEB is implemented using PySpark.

###  3.2. MicroSynth

MicroSynth differs from EB in the sense that it adopts a truncated square
loss, so the solution would typically be found at the corner of the constraint
set and lead to sparse weights. Besides, MicroSynth prevents outlying weights
by punishing more on large weights than EB (Robbins et al., 2017). The
optimization problem of MicroSynth is much more challenging than EB because it
involves constraints that prevent negative weight. For simplicity, we add a
constant into covariate vector
$X_{i}^{\prime}=[1,X_{i1},X_{i2},\cdots,X_{id}]^{\top}$, and the corresponding
treatment mean
$\tilde{X}^{\prime}=\frac{1}{N_{1}}\sum_{T_{i}=1}X_{i}^{\prime}$. Let
$A=[X_{1}^{\prime},X_{2}^{\prime},\cdots,X_{N_{0}}^{\prime}]$ be a matrix of
covariates for all users in the control group and
$\bm{w}=[w_{1},w_{2},\cdots,w_{N_{0}}]^{\top}$ be the vector of weights, then
the equality constraints in Equation 5 can be expressed as
$A\bm{w}=\tilde{X}^{\prime}$. We apply Lagrange multipliers
$\bm{\xi}\in\mathbb{R}^{d+1}$ like in EB and the primal problem is

(11)
$\textrm{min}\bigg{\\{}\mathcal{L}^{\textrm{MicroSynth}}=\frac{1}{2}||\bm{w}-\mathbf{1}||_{2}^{2}+\psi_{+}(\bm{w})+\bm{\xi}^{\top}(A\bm{w}-\tilde{X}^{\prime})\bigg{\\}},$

where $\psi_{+}(\bm{w})=\infty$ if any $w_{i}<0$ and 0 otherwise. This is a
high dimensional quadratic programming problem with both equality and
inequality constraints. The number of parameters can be at scale of hundreds
of millions, and it is infeasible to scale up traditional solvers such as
interior point, projected gradient methods, active set, etc. We have found
that dual ascent method is well suited for large scale distributed convex
optimization problems. Another option that we have considered is Augmented
Lagrangian methods like ADMM (Boyd et al., 2011), however we find that dual
ascent is computationally more efficient.

The dual ascent method works in an iterative fashion. At each iteration we use
gradient ascent to update the Lagrange dual variable, the optimal primal
variable can then be recovered by minimizing the primal problem. The dual
ascent method utilizes the fact that the Lagrange dual problem is always
concave, unconstrained, and relatively low dimensional, so we are able to use
gradient based methods in a distributed manner to solve the dual problem. Let
$f(\bm{w})=\frac{1}{2}||\bm{w}-\mathbf{1}||_{2}^{2}+\psi_{+}(\bm{w})$, the
dual problem of Equation 11 is

(12)
$\max_{\bm{\xi}}\bigg{\\{}\inf_{\bm{w}}\mathcal{L}^{\textrm{MicroSynth}}=-f^{*}(-A^{\top}\bm{\xi})-\tilde{X}^{\prime\top}\bm{\xi}\bigg{\\}},$

where $f^{*}$ is the convex conjugate of $f$. Because $f({\bm{w}})$ is
strictly convex, $f^{*}$ is differentiable. Using properties of conjugates,
the derivative of Equation 12 is $A{\bm{w}}^{+}-\tilde{X}^{\prime}$ where
$\bm{w}^{+}={\mathrm{argmin}_{\bm{w}}}~{}\mathcal{L}^{\textrm{MicroSynth}}(\bm{w},\bm{\xi})$.
Therefore the dual ascent method starts with an initial $\bm{\xi}^{(0)}$ and
updates $\bm{w}^{(k)},\bm{\xi}^{(k)}$ alternatively for $K$ iterations:

(13) $\displaystyle\bm{w}^{(k+1)}$
$\displaystyle=\underset{\bm{w}}{\mathrm{argmin}}~{}\bigg{\\{}\frac{1}{2}\|\bm{w}-\mathbf{1}\|_{2}^{2}+\psi_{+}(\bm{w})+\bm{\xi}^{(k)\top}(A\bm{w}-\tilde{X}^{\prime})\bigg{\\}},$
(14) $\displaystyle\bm{\xi}^{(k+1)}$
$\displaystyle=\bm{\xi}^{(k)}+\alpha^{(k)}(A\bm{w}^{(k+1)}-\tilde{X}^{\prime}).$

By calculus and projection, the problem in Equation 13 has a close-form
solution:

(17) $\displaystyle w_{i}^{(k+1)}$
$\displaystyle=\left\\{\begin{array}[]{ll}1-X_{i}^{\top}\bm{\xi}^{(k)},&\mbox{
if }1-X_{i}^{\top}\bm{\xi}^{(k)}>0\\\ 0,&\mbox{ otherwise.}\end{array}\right.$

Notice that this solution can be easily calculated using MapReduce because the
weight $w_{i}$ only depends on the covariates of user $i$. In practice we add
momentum in the dual gradient ascent step to speed up the convergence. The
full detail of DistMS is described in Algorithm 2.

Algorithm 2 Distributed MicroSynth (DistMS)

1:Initialize $\bm{\xi}^{(0)}\leftarrow\mathbf{0}$, momentum
$\bm{v}^{(0)}\leftarrow\mathbf{0}$, learning rate $\alpha=0.01,\beta^{(0)}=1$.

2:for $k=1,\ldots,K$ do $\triangleright$ $K$: the number of iterations

3:
$w_{i}^{(k)}\leftarrow\left\\{\begin{array}[]{ll}1-X_{i}^{\top}\bm{\xi}^{(k-1)},\mbox{
if }1-X_{i}^{\top}\bm{\xi}^{(k-1)}>0\\\ 0,\mbox{
otherwise.}\end{array}\right.$ $\triangleright$ Distributed compute weights
using Map

4:
$\bm{v}^{(k)}\leftarrow\bm{\xi}^{(k-1)}+\alpha(A\bm{w}^{(k)}-\tilde{X}^{\prime})$
$\triangleright$ Dual ascent step, MapReduce is used to compute
$A\bm{w}^{(k)}$

5: $\beta^{(k)}\leftarrow\frac{1+\sqrt{1+4(\beta^{(k-1)})^{2}}}{2}$
$\triangleright$ Increase the 2nd learning rate

6:
$\bm{\xi}^{(k)}\leftarrow\bm{v}^{(k)}+\Big{(}\frac{\beta^{(k-1)}-1}{\beta^{(k)}}\Big{)}(\bm{v}^{(k)}-\bm{v}^{(k-1)})$
$\triangleright$ Momentum step

7:end for

8:Compute final weights $w_{i}^{\textrm{MicroSynth}}=w_{i}^{(K)}$ for all
control users. $\triangleright$ MapReduce

Similar to DistEB, DistMS loops through the whole data set twice in each
iteration and the overall time complexity is $O(NdK)$.

###  3.3. Comparison with other implementations

Our implementations exploit the Lagrange dual problem like most of the
existing implementations of balancing approaches. The difference lies in the
way we solve the dual problem. ebal uses a Levenberg-Marquardt scheme to
iteratively solve the problem (Hainmueller, 2022). The drawback is that it
requires calculation of the Hessian matrix, which is infeasible under high
dimensional settings. EC directly uses scipy.optimize.fsolve to solve the dual
problem (Wang et al., 2019). Under the hood, scipy.optimize.fsolve calculates
the Jacobian matrix by a forward-difference approximation, then uses a
modified Powell method to iteratively find the solution. This approach is
faster than the Hessian matrix calculation. Our method is even faster in the
sense that the gradient of the dual variable is directly computed from the
data instead of a forward-difference step with a lot of extra function
evaluations. Besides, our method is easier to implement in a distributed
fashion. We compare the scalability of our algorithms with other
implementations and find that our algorithms are faster than the existing in-
memory implementations of EB and MicroSynth, and competitive with
mathematically simpler methods such as IPW. The detailed results are described
in Section 4.5.

###  3.4. Distributed vs in-memory

Although our main focus is to design distributed algorithms for EB and
MicroSynth, the algorithms can also be used serially and implemented in-
memory. We implement in-memory version of both DistEB and DistMS in Python,
utilizing NumPy’s fast vectorized operations. When the data set is small and
fits into a single-node machine, the in-memory version is very efficient
because in-memory operations are much faster at different orders of magnitude
than data shuffling between machines. However, machines with Terabytes of
memories are not always available, in which case we need to distribute the
workload on a cluster of workers. Moreover, when the data size is large,
loading data into memory will be slow and take up a large portion of the total
running time. We provide critical code snippets and other details for both the
in-memory and Spark implementation of EB and MicroSynth in supplementary
materials (Section C).

##  4\. Evaluation

###  4.1. Data Generating Process

We use Monte Carlo simulations to evaluate the performance of different
weighting methods on semi-synthetic data. We compare DistEB with DistMS, IPW
based on logistic regression,222Other IPW procedures, such as machine-learning
based approach like boosted regressions (Mccaffrey et al., 2005) can also be
applied to calculate the propensity score. However, with model selection and
hyper-parameter tuning, these machine-learning based approach could be very
expensive in big data. Therefore, we do not discuss the performance of these
methods. CBPS,333We use Meta’s python package balance to implement CBPS. We
focus on over-identified parametric CBPS which simultaneously fits a
propensity score model and generates balancing weights. There is also a just-
identified version of CBPS which only generates balancing weights without
maximizing the treatment selection prediction. We also apply this method in
simulation and find the results are very similar to the over-identified CBPS.
We do not implement non-parametric CBPS which involves in very intensive
calculation. and the combination of these methods and a doubly robust (DR)
process in Equation 4 by fitting the outcome variables with simple Ordinary
Least Square regressions. To make the data generating process as realistic as
possible, instead of using fully simulated data with normal distribution, we
base our simulations on semi-synthetic data so that our simulations reflect
9-week real user data on the Snapchat platform. We define week 1 to week 8 as
pretreatment period and week 9 as post-treatment period.

We simulate data with 487 baseline covariates. It includes 456 continuous
covariates which represents weekly user engagement of 57 major engagement
metrics from week 1 to week 8. Many of the user engagement metrics are skewed
and similar to the zero-inflated negative binomial distribution. We also have
31 binary covariates which refer to major user attributes in week 8. We use
three representative user engagement metrics: active day, app open, time spent
in week 9 as the outcome variables. The self-selected treatment assignment
metric is users’ OS types. Users are in the treated group if OS type is iOS
and control group if OS type is Android.

We assume that all 487 covariates could influence treatment assignment and
outcome variables. We focus on three specifications of outcome models
$\widehat{g_{j}}$ and selection model $\widehat{f}$: linear model, linear
model with interactions, and random forest model. The detailed data generating
steps are in supplementary materials (Algorithm 3). In the linear
specification of outcome models and selection model, we achieve the correct
specification. No regularization parameter is added to the IPW method to
prevent misspecification bias in this specification. In the other two
specifications, we would have misspecification problems by applying DistEB,
DistMS, IPW and CBPS and only including first moment. We tune regularizaion
parameter to achieve the best covariate balance (see supplementary materials
Section D for details). Hyperparameter tuning improves the performance of IPW
in these two specifications. The purpose is to compare the performance of
several weighting methods when both treatment selection model and outcome
models are misspecified, since we never know the correct model in real-world
practice.

We evaluate the performance of each method in three different sample size:
10K, 100K, and 1 Million to show the performance gain as sample size
increases. We generate 100 datasets for each sample size by random sampling.

###  4.2. Meta Metrics Evaluation

We compare different methods and sample sizes in bias, variance, balance, and
weight stability.

We measure the bias of each outcome variable $Y_{j}$ by Absolute Mean Bias:
$AMB_{j}=\frac{1}{100}|\sum_{s=1}^{100}\widehat{\delta_{js}}-\delta_{j}|$,
where
$\widehat{\delta_{js}}=\frac{\widehat{\tau_{js,PATT}}}{\overline{Y_{js,T_{i}=1}}-\widehat{\tau_{js,PATT}}}$.
Since different outcome variables vary in magnitude, instead of PATT, we focus
on the percentage difference between the treated group and the re-weighted
control group. The ground truth percentage change $\delta_{j}$ is always 0
since we do not introduce any additional treatment effect.

We measure the variance by Standard Deviation:
$SD_{j}=sd(\widehat{\delta_{js}})$. It is noteworthy that another popular
evaluation meta metric Root Mean Square Error (RMSE) can be decomposed to bias
and variance as: $RMSE_{j}^{2}=AMB_{j}^{2}+SD_{j}^{2}$. We do not report it
separately.

We can further understand bias and variance by covariate balance. Balance
means similarity between the treated group and the synthetic control group in
pretreatment covariates. We use several popular balance metrics for weighting
methods in literature ((Franklin et al., 2014; Imai et al., 2007)) as follows:
standardized mean difference (SMD), variance Ratio (VR), Overlapping
coefficient (OVL), Mahalanobis Balance (MB), and Kolmogorov-Smirnov distance
(KS). The detailed definition of these metrics are in supplementary materials
(Section B). We summarize the average balance across all dimensions. Among
these balance metrics, SMD of original metrics ($X_{j}$) and MB measure the
balance at mean. SMD of quadratic terms ($X_{j}^{2}$), SMD of interaction
terms ($X_{j}X_{k}$), and VR measure the balance at second moments. OVL and KS
measure the balance at the whole distribution.

Another diagnostic metric group is weight stability. Low stability or high
variability in weights often means a few units receive very large weights. It
leads to an unstable PATT estimator. Follow (Chattopadhyay et al., 2020), we
focus on the following meta metrics for weight stability measurement: Standard
Deviation of Normalized Weights (SD), extremity of weights (measured by the
maximum (MAX) and 99th percentile (P99) of weights), and Effective Sample Size
(ESS). The detailed definition of these metrics are in supplementary materials
(Section B). A low level of ESS indicates the presence of a few units with
extreme weights. Therefore, higher value of this metric indicates better
weight stability. In our simulation we use $ESS$ divided by the total number
of units in the control group to report the ratio of EES. This ratio can be
used as a standalone diagnostic in practice (Chattopadhyay et al., 2020).

###  4.3. Bias, Variance, Balance, Weight Stability in Different Data
Generating Processes

We evaluate the performance of different weighting methods across three data
generation scenarios, with a sample size of 100K as the primary focus. Figure
1 shows that no matter in correctly specified model (Linear) or incorrectly
specified models (Linear with Interaction and Random Forest), DistEB and
DistMS always have much smaller bias and smaller or similar variance than IPW
and CBPS. When both selection and outcome models are misspecfied, DistEB and
DistMS still have very small bias. It means by achieving covariate balancing,
results are robust to misspecification. In real world, we always have more or
less misspecification problems since we never know the ground truth data
generating process. DistEB and DistMS have similar performance in bias and
variance to IPW+DR and CBPS+DR. DistMS has slightly smaller variance than
DistEB in all cases. DR significantly reduces the bias and variance of IPW and
CBPS, but fails to enhance the performance of DistEB and DistMS. It aligns
with the theory that EB and MicroSynth are implicitly doubly robust (Zhao and
Percival, 2017). CBPS has worse bias but better variance than IPW in our data.
It is consistent with other literature that in both correctly specified models
and incorrectly specified models, compared to IPW, CBPS could have worse
performance in bias but better performance in variance (Zhao and Percival,
2017).

The good performance of DistEB and DistMS can be attributed to their ability
to balance covariates. It is noteworthy that we do not include higher-order
moments in calculating the weights. The two balancing methods not only have
good performance at mean, but also have good performance at second moment and
the whole distribution. In every instance, the balance of DistEB and DistMS
surpasses the recommended threshold according to literature.444We put the
recommended threshold in literature in the supplementary materials. See
Section B for details. IPW frequently falls below the threshold. Our
simulations reveal that even a small imbalance can result in substantial bias.
DistEB and DistMS can mitigate bias by enhancing balance, while CBPS aims to
enhance balance, but fails to do so in our data.

The stability of the PATT estimator is related to weight stability. We find
that even though CBPS has bad performance in bias and balance, it is the one
has the best weight stability. Theoretically, EB is more likely to produce
extreme weights than MicroSynth (Robbins et al., 2017). Our simulations
confirms this statement. The performance of IPW is volatile in different
simulations and is worse than DistMS and CBPS.

It’s important to note that the IPW method is costly for both model selection
and hyperparameter tuning, and does not guarantee the selection of a model
with the lowest bias in both expectation and specific finite samples. The
reason is that when the propensity score model is correctly specified, IPW
only balances covariates and achieves no bias in expectations but not in any
particular finite sample. Our findings show that when the data generation
process is linear, tuning hyperparameters for optimal covariate balance in
each simulation does not select the correctly specified model or the model
with the lowest bias.

Figure 1. Comparison of weighting methods in bias, variance, balance, and
weight stability in three data generating processes (DGP). Each DGP creates
100 simulations with sample size as 100,000. DistEB and DistMS consistently
outperform IPW and CBPS. The upper-left panel shows the comparison in bias
measured by absolute mean bias across all simulations. The upper-right panel
shows variance measured by standard deviation. The lower-left and lower-right
panel show balance measured by seven balance metrics and weight stability
measured by four metrics, respectively. The dotted line is 1% for absolute
mean bias and standard deviation and recommended threshold in literature for
balance. Y-axis is sqrt transformed for better visualization.

###  4.4. Bias, Variance, Balance, Weight Stability by Sample Size

We evaluate the performance of various methods across three sample sizes: 10K,
100K, and 1 Million, with a focus on data generated by the random forest
method. Figure 2 shows that bias and balance greatly improve from 10K to 100K
but are quite similar in most cases if we compare 100K and 1 Million. Variance
and some weight stability metrics always decreases as sample size increases.
It means large sample size is necessary for us to get stable estimators and
detect small change.

Figure 2. Comparison of weighting methods in bias, variance, balance, and
weight stability for three sample sizes. With larger sample sizes, all methods
performs better. The upper and the middle panel show the comparison in bias
and variance, respectively. The lower-left and the lower-right panel show
balance measured by seven balance metrics and the weight stability measured by
four metrics, respectively. For each sample size, the results are derived
using 100 simulations with data generating process as random forest. The
dotted line is 1% for absolute mean bias and standard deviation and the
recommended threshold in literature for balance. Y-axis is sqrt transformed
for better visualization. We do not calculate Mean (SMD interaction) for
1Million sample size because the computation is expensive.

###  4.5. Running Time and Scalability

We evaluate the in-memory and distributed version of DistEB and DistMS on
Google Cloud Platform general purpose machines, details can be found in
supplementary materials.

We first evaluate the in-memory version on a single-node machine and compare
the performance of our method with existing open source packages including EB
and MicroSynth by Google’s EC package (Wang et al., 2019), CBPS by Meta’s
balance package, and IPW at different scales of data points (10K, 100K,
1Million). As shown in Table 1, compared to the existing packages of EB and
MicroSynth, our in-memory version has similar or better performance in terms
of running time. In general, MicroSynth outperforms EB because of faster
convergence, thus requires fewer iterations. Both MicroSynth and EB
significantly outperforms CBPS. IPW with L2 regularization has similar
performance to MicroSynth. However, to achive the best performance, IPW
requires parameter tuning, which is not included in the running time.

Table 1. Running Time in minutes on a Single-node Machine Data Size
---
| 10K | 100K | 1Million
DistEB (this work) | 0.248 | 1.375 | 5.597
EB (EC) | 2.259 | 6.066 | 124.317
DistMS (this work) | 0.05 | 0.142 | 0.995
MicroSynth (EC) | 2.342 | 0.125 | 0.99
CBPS | 6.411 | 44.628 | 433
IPW | 0.015 | 0.213 | 0.783

For the distributed version of DistMS, we run the algorithm on a 10 Million
dataset over clusters with different number of workers to demonstrate the
scalability. As shown in Figure 3, running time decreases as the number of
workers increases, and stabilizes as the number of parallel tasks in Spark
gets closer to the number of file partitions. Distributed computing also
avoids the long data loading time on a single machine.

Figure 3. Running Time of DistMS in Spark on a 10 Million dataset

##  5\. Applications at Snap Inc.

At Snap Inc., We combine the balancing approaches and the synthetic control
framework with multiple treated units (Abadie, 2021) and design an end-to-end
workflow for observational study analysis, as shown in Figure 4. With the
synthetic control framework, we balance the pretreatment time-series of
multiple engagement metrics and user attributes to mitigate the bias caused by
unmeasured confounders (Abadie et al., 2010). We construct a synthetic control
unit with millions of users in the control donor pool. The synthetic control
unit serves as the counterfactual of the treated units. Considering the
performance in variance, weight stability, and running time, DistMS
outperforms DistEB. We productionize DistMS in the workflow. Our tool takes
the user identifiers in the treated group, control donor pool and the
treatment start date as inputs, then uses DistMS to generate a set of weights
for users in the control donor pool. The re-weighted control group is directly
comparable with the treated group. In the end, we provide a summary report
containing the cumulative differences and time series. Following the best
practices of the synthetic control method (Abadie, 2021), we run balance check
on all balancing variables including pretreatment user engagement metrics and
user attributes such as country, age, gender etc, which is analogical to the
concept of sample ratio mismatch (Fabijan et al., 2019) from A/B testing. We
leave one week out of matching period as the validation period, which mimics
the AA period in A/B testing. Confidence intervals are calculated using
bootstrap method. Users in familiarity with traditional A/B testing platform
can easily interpret the results of observational studies, given the
similarity between the reporting interfaces, as demonstrated in Figure 5 and
Figure 6.

This end-to-end workflow has been applied to many cases at Snap Inc. Engineers
can use the tool to analyze the impact on user engagement of some events or
incidents that happen naturally, such as network errors that only happen in a
certain area. Data scientists can also design quasi-experiments, such as
launching a marketing campaign for Snapchat in a college and measuring the
impact with synthetic control. Moreover, in randomized experiments, app
performance metrics such as app open latency could suffer from self-selection
biases, because the metric is only reported when users open the app, as
pointed out in (Xie et al., 2020). Balancing approaches can be directly
applied to mitigate the biases by adjusting for relevant covariates. Other
companies also face these scenarios and can apply this system to their
analysis.

Figure 4. End-to-end workflow for observational study analysis at Snap Inc.
Figure 5. Summary Report Figure 6. Time Series Report: We have 12 pretreatment
weeks, 1 validation week, and 7 post-treatment weeks in this example. The red
dotted line separates pretreatment period and validation period. The red
dashed line separates validation period and post-treatment period.

##  6\. Discussion

The validity of the end-to-end system relies on certain assumptions.

Firstly, similar to most causal inference methods with observational data, we
assume no unmeasured confounders and overlap. With larger datasets, we can
balance numerous confounders, mitigating biases caused by unmeasured
confounders. In practice, we focus on pre-treatment time-series of key
engagement metrics and user attributes, balancing only the first moment by
default. In some cases, balancing more pre-treatment covariates can help
reduce bias caused by unobserved confounders. The validation period before
treatment allows us to check for strong evidence of unobserved confounders.

Secondly, we have the assumption of synthetic control framework. This
assumption asserts that a good pre-treatment fit acts as a proxy for a good
accounting of all observed and unobserved factors. Matching pre-intervention
outcome time series helps control for heterogeneous responses to multiple
unobserved factors. With this assumption, only unobserved confounders in the
post-treatment period, which cannot be accurately predicted by pre-treatment
information, can bias our results.

Lastly, our system relies on the data generating assumption underlying the
balancing approach. This assumption requires a linear true data generating
process for either the selection model or the outcome model. We find that this
assumption can be loosen in practice. Our simulations demonstrate that even
when the true data generating process is nonlinear, our system still provides
estimates with smaller bias compared to other weighting methods. Nevertheless,
including higher moments may be helpful in reducing misspecification bias in
certain cases (Huang et al., 2021). Incorporating a large number of balancing
covariates and moments poses two challenges. Firstly, finding a solution that
balances all covariates and moments becomes more difficult as more constraints
are added to the optimization procedure. Secondly, including many balancing
constraints increases computational costs. Therefore, applying feature
selection or dimension reduction methods before the balance procedure can be
helpful, and this aspect remains open for future improvement.

In summary, our findings demonstrate that the balancing approach, which
directly balances confounders, outperforms the modeling approach by producing
a causal effect estimator with low bias and variance. This approach is also
more efficient, as it does not require modeling multiple outcome variables and
adjusting hyperparameters, resulting in a doubly robust solution with low
computational cost. A large sample size also reduces estimation variance. To
further improve the solution, we propose scalable solvers that have faster in-
memory computation than existing solvers and can be easily integrated with
distributed computing frameworks. Our application combines the balancing
approach with a synthetic control framework to address unmeasured confounders.
This end-to-end workflow is successfully implemented at Snap Inc. to estimate
causal effects using large-scale observational data.

## References

  * (1)
  * Abadie (2021) Alberto Abadie. 2021\.  Using Synthetic Controls: Feasibility, Data Requirements, and Methodological Aspects.  _Journal of Economic Literature_ 59, 2 (June 2021), 391–425.  https://doi.org/10.1257/jel.20191450
  * Abadie et al. (2010) Alberto Abadie, Alexis Diamond, and Jens Hainmueller. 2010\.  Synthetic Control Methods for Comparative Case Studies: Estimating the Effect of California’s Tobacco Control Program.  _J. Amer. Statist. Assoc._ 105, 490 (2010), 493–505.  https://doi.org/10.1198/jasa.2009.ap08746 arXiv:https://doi.org/10.1198/jasa.2009.ap08746
  * Abadie and Imbens (2006) Alberto Abadie and Guido Imbens. 2006.  Large Sample Properties of Matching Estimators for Average Treatment Effects.  _Econometrica_ 74, 1 (2006), 235–267.  https://EconPapers.repec.org/RePEc:ecm:emetrp:v:74:y:2006:i:1:p:235-267
  * Battocchi et al. (2019) Keith Battocchi, Eleanor Dillon, Maggie Hei, Greg Lewis, Paul Oka, Miruna Oprescu, and Vasilis Syrgkanis. 2019.  EconML: A Python Package for ML-Based Heterogeneous Treatment Effects Estimation.  https://github.com/microsoft/EconML.  Version 0.x.
  * Bluhm and Cutura (2022) Benjamin Bluhm and Jannic Alexander Cutura. 2022.  Econometrics at Scale: Spark up Big Data in Economics.  _Journal of Data Science_ 20, 3 (2022), 413–436.  https://doi.org/10.6339/22-JDS1035
  * Boyd et al. (2011) Stephen P. Boyd, Neal Parikh, Eric Chu, Borja Peleato, and Jonathan Eckstein. 2011.  Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers.  _Foundations and Trends in Machine Learning_ 3, 1 (2011), 1–122.  http://dblp.uni-trier.de/db/journals/ftml/ftml3.html#BoydPCPE11
  * Brodersen et al. (2014) Kay H. Brodersen, Fabian Gallusser, Jim Koehler, Nicolas Remy, and Steven L. Scott. 2014.  Inferring causal impact using Bayesian structural time-series models.  _Annals of Applied Statistics_ 9 (2014), 247–274.  https://research.google/pubs/pub41854/
  * Chan et al. (2016) Kwun Chuen Gary Chan, Sheung Chi Phillip Yam, and Zheng Zhang. 2016.  Globally efficient non-parametric inference of average treatment effects by empirical balancing calibration weighting.  _Journal of the Royal Statistical Society. Series B (Statistical Methodology)_ 78, 3 (2016), 673–700.  http://www.jstor.org/stable/24775357
  * Chattopadhyay et al. (2020) Ambarish Chattopadhyay, Christopher Hase, and José Zubizarreta. 2020.  Balancing vs modeling approaches to weighting in practice.  _Statistics in Medicine_ 39 (09 2020).  https://doi.org/10.1002/sim.8659
  * Chen et al. (2020) Huigang Chen, Totte Harinen, Jeong-Yoon Lee, Mike Yung, and Zhenyu Zhao. 2020.  Causalml: Python package for causal machine learning.  _arXiv preprint arXiv:2002.11631_ (2020).
  * Chen et al. (2022) Rui Chen, Guanhua Chen, and Menggang Yu. 2022.  Entropy Balancing for Causal Generalization with Target Sample Summary Information.  (Jan 2022).
  * Fabijan et al. (2019) Aleksander Fabijan, Jayant Gupchup, Somit Gupta, Jeff Omhover, Wen Qin, Lukas Vermeer, and Pavel Dmitriev. 2019.  Diagnosing sample ratio mismatch in online controlled experiments.  _Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_ (2019).  https://doi.org/10.1145/3292500.3330722
  * Franklin et al. (2014) Jessica Franklin, Jeremy Rassen, Diana Ackermann, Dorothee Bartels, and Sebastian Schneeweiss. 2014\.  Metrics for covariate balance in cohort studies of causal effects.  _Statistics in medicine_ 33 (05 2014).  https://doi.org/10.1002/sim.6058
  * Hainmueller (2012) Jens Hainmueller. 2012\.  Entropy Balancing for Causal Effects: A Multivariate Reweighting Method to Produce Balanced Samples in Observational Studies.  _Political Analysis_ 20, 1 (2012), 25–46.  https://doi.org/10.1093/pan/mpr025
  * Hainmueller (2022) Jens Hainmueller. 2022\.  _ebal: Entropy Reweighting to Create Balanced Samples_.  https://CRAN.R-project.org/package=ebal R package version 0.1-8.
  * Hainmueller and Xu (2013) Jens Hainmueller and Yiqing Xu. 2013.  ebalance: A Stata Package for Entropy Balancing.  _Journal of Statistical Software_ 54, 7 (2013), 1–18.  https://doi.org/10.18637/jss.v054.i07
  * Hazlett (2020) Chad Hazlett. 2020\.  KERNEL BALANCING: A FLEXIBLE NON-PARAMETRIC WEIGHTING PROCEDURE FOR ESTIMATING CAUSAL EFFECTS.  _Statistica Sinica_ 30, 3 (2020), 1155–1189.  https://www.jstor.org/stable/26968924
  * Hirano and Imbens (2001) Keisuke Hirano and Guido W. Imbens. 2001.  Estimation of causal effects using propensity score weighting: An application to data on right heart catheterization.  _Health Services and Outcomes Research Methodology_ 2, 3-4 (2001), 259–278.  https://doi.org/10.1023/A:1020371312283
  * Hirano et al. (2003) Keisuke Hirano, Guido Imbens, and Geert Ridder. 2003\.  Efficient Estimation of Average Treatment Effects Using the Estimated Propensity Score.  _Econometrica_ 71 (02 2003), 1161–1189.  https://doi.org/10.1111/1468-0262.00442
  * Huang et al. (2021) Melody Y. Huang, Brian G. Vegetabile, Lane F. Burgette, Claude Setodji, and Beth Ann Griffin. 2021\.  Balancing Higher Moments Matters for Causal Estimation: Further Context for the Results of Setodji et al. (2017).  https://doi.org/10.48550/ARXIV.2107.03922
  * Huling and Mak (2020) Jared D. Huling and Simon Mak. 2020.  Energy Balancing of Covariate Distributions.  https://doi.org/10.48550/ARXIV.2004.13962
  * Imai et al. (2007) Kosuke Imai, Gary King, and Elizabeth Stuart. 2007.  Misunderstandings Between Experimentalists and Observationalists About Causal Inference.  _Journal of the Royal Statistical Society, Series A_ 171 (12 2007).
  * Imai et al. (2008) Kosuke Imai, Gary King, and Elizabeth Stuart. 2008.  Misunderstandings Among Experimentalists and Observationalists about Causal Inference.  _Journal of the Royal Statistical Society, Series A_ 171, part 2 (2008), 481–502.
  * Imai and Ratkovic (2014) Kosuke Imai and Marc Ratkovic. 2014.  Covariate balancing propensity score.  _Journal of the Royal Statistical Society: Series B (Statistical Methodology)_ 76, 1 (2014), 243–263.  https://doi.org/10.1111/rssb.12027 arXiv:https://rss.onlinelibrary.wiley.com/doi/pdf/10.1111/rssb.12027
  * Josey et al. (2021) Kevin Josey, Seth Berkowitz, Debashis Ghosh, and Sridharan Raghavan. 2021.  Transporting experimental results with entropy balancing.  _Statistics in Medicine_ 40 (05 2021).  https://doi.org/10.1002/sim.9031
  * Kang and Schafer (2007) Joseph D. Y. Kang and Joseph L. Schafer. 2007.  Demystifying Double Robustness: A Comparison of Alternative Strategies for Estimating a Population Mean from Incomplete Data.  _Statist. Sci._ 22, 4 (nov 2007).  https://doi.org/10.1214/07-sts227
  * Lunceford and Davidian (2004) Jared K Lunceford and Marie Davidian. 2004.  Stratification and weighting via the propensity score in estimation of causal treatment effects: a comparative study.  _Statistics in Medicine_ 23 (2004).
  * Markoulidakis et al. (2022) Andreas Markoulidakis, Khadijeh Taiyari, Peter Holmans, Philip Pallmann, Monica Busse, Mark Godley, and Beth Ann Griffin. 2022\.  A tutorial comparing different covariate balancing methods with an application evaluating the causal effects of substance use treatment programs for adolescents.  _Health Services and Outcomes Research Methodology_ (05 2022).  https://doi.org/10.1007/s10742-022-00280-0
  * Mccaffrey et al. (2005) Daniel Mccaffrey, Greg Ridgeway, and Andrew Morral. 2005\.  Propensity Score Estimation With Boosted Regression for Evaluating Causal Effects in Observational Studies.  _Psychological methods_ 9 (01 2005), 403–25.  https://doi.org/10.1037/1082-989X.9.4.403
  * Ouattara et al. (2021) Aurélien Ouattara, Matthieu Bulté, Wan-Ju Lin, Philipp Scholl, Benedikt Veit, Christos Ziakas, Florian Felice, Julien Virlogeux, and George Dikos. 2021.  Scalable Econometrics on Big Data – The Logistic Regression on Spark.  _arXiv e-prints_ , Article arXiv:2106.10341 (June 2021), arXiv:2106.10341 pages.  https://doi.org/10.48550/arXiv.2106.10341 arXiv:2106.10341 [stat.CO]
  * Qian (1999) Ning Qian. 1999\.  On the momentum term in gradient descent learning algorithms.  _Neural Networks_ 12, 1 (1999), 145–151.  https://doi.org/10.1016/S0893-6080(98)00116-6
  * Robbins and Davenport (2021) Michael W. Robbins and Steven Davenport. 2021.  microsynth: Synthetic Control Methods for Disaggregated and Micro-Level Data in R.  _Journal of Statistical Software_ 97, 2 (2021), 1–31.  https://doi.org/10.18637/jss.v097.i02
  * Robbins et al. (2017) Michael W. Robbins, Jessica Saunders, and Beau Kilmer. 2017\.  A Framework for Synthetic Control Methods With High-Dimensional, Micro-Level Data: Evaluating a Neighborhood-Specific Crime Intervention.  _J. Amer. Statist. Assoc._ 112, 517 (2017), 109–126.  https://doi.org/10.1080/01621459.2016.1213634 arXiv:https://doi.org/10.1080/01621459.2016.1213634
  * Robins and Rotnitzky (1995) James M. Robins and Andrea Rotnitzky. 1995.  Semiparametric Efficiency in Multivariate Regression Models with Missing Data.  _J. Amer. Statist. Assoc._ 90, 429 (1995), 122–129.  http://www.jstor.org/stable/2291135
  * Robins et al. (1994) James M. Robins, Andrea Rotnitzky, and Lue Ping Zhao. 1994\.  Estimation of Regression Coefficients When Some Regressors are not Always Observed.  _J. Amer. Statist. Assoc._ 89, 427 (1994), 846–866.  https://doi.org/10.1080/01621459.1994.10476818 arXiv:https://doi.org/10.1080/01621459.1994.10476818
  * Rosenbaum and Rubin (1983) Paul R. Rosenbaum and Donald B. Rubin. 1983.  The central role of the propensity score in observational studies for causal effects.  _Biometrika_ 70, 1 (04 1983), 41–55.  https://doi.org/10.1093/biomet/70.1.41 arXiv:https://academic.oup.com/biomet/article-pdf/70/1/41/662954/70-1-41.pdf
  * Rosenbaum and Rubin (1984) Paul R. Rosenbaum and Donald B. Rubin. 1984.  Reducing Bias in Observational Studies Using Subclassification on the Propensity Score.  _J. Amer. Statist. Assoc._ 79, 387 (1984), 516–524.  https://doi.org/10.1080/01621459.1984.10478078 arXiv:https://www.tandfonline.com/doi/pdf/10.1080/01621459.1984.10478078
  * Rosenbaum and Rubin (1985) Paul R. Rosenbaum and Donald B. Rubin. 1985.  Constructing a Control Group Using Multivariate Matched Sampling Methods That Incorporate the Propensity Score.  _The American Statistician_ 39, 1 (1985), 33–38.  https://doi.org/10.1080/00031305.1985.10479383 arXiv:https://www.tandfonline.com/doi/pdf/10.1080/00031305.1985.10479383
  * Rubin (2001) Donald Rubin. 2001\.  Using Propensity Scores to Help Design Observational Studies: Application to the Tobacco Litigation.  _Health Services and Outcomes Research Methodology_ 2 (12 2001), 169–188.  https://doi.org/10.1023/A:1020363010465
  * Rubin (1980) Donald B. Rubin. 1980\.  Randomization Analysis of Experimental Data: The Fisher Randomization Test Comment.  _J. Amer. Statist. Assoc._ 75, 371 (1980), 591–593.  http://www.jstor.org/stable/2287653
  * Rubin (2008) Donald B. Rubin. 2008\.  For objective causal inference, design trumps analysis.  _The Annals of Applied Statistics_ 2, 3 (sep 2008).  https://doi.org/10.1214/08-aoas187
  * Sharma and Kiciman (2020) Amit Sharma and Emre Kiciman. 2020.  DoWhy: An End-to-End Library for Causal Inference.  _arXiv preprint arXiv:2011.04216_ (2020).
  * Stuart (2010) Elizabeth A. Stuart. 2010\.  Matching Methods for Causal Inference: A Review and a Look Forward.  _Statist. Sci._ 25, 1 (2010), 1–21.  http://www.jstor.org/stable/41058994
  * Wang et al. (2019) Xiaojing Wang, Jingang Miao, and Yunting Sun. 2019.  A Python Library For Empirical Calibration.  _arXiv preprint arXiv:1906.11920_ (2019).
  * Wong et al. (2021) Jeffrey Wong, Eskil Forsell, Randall Lewis, Tobias Mao, and Matthew Wardrop. 2021.  You Only Compress Once: Optimal Data Compression for Estimating Linear Models.  _arXiv preprint arXiv:2102.11297_ (2021).
  * Wong (2020) Jeffrey C Wong. 2020\.  Computational causal inference.  _arXiv preprint arXiv:2007.10979_ (2020).
  * Xie et al. (2020) Yuxiang Xie, Meng Xu, Evan Chow, and Xiaolin Shi. 2020\.  How to Measure Your App: A Couple of Pitfalls and Remedies in Measuring App Performance in Online Controlled Experiments.  _Proceedings of the 14th ACM International Conference on Web Search and Data Mining_ (2020).
  * Zhao (2019) Qingyuan Zhao. 2019\.  Covariate balancing propensity score by tailored loss functions.  _The Annals of Statistics_ 47, 2 (2019), 965 – 993.  https://doi.org/10.1214/18-AOS1698
  * Zhao and Percival (2017) Qingyuan Zhao and Daniel Percival. 2017.  Entropy Balancing is Doubly Robust.  _Journal of Causal Inference_ 5, 1 (2017), 20160010\.
  * Zubizarreta (2015) José Zubizarreta. 2015\.  Stable Weights that Balance Covariates for Estimation With Incomplete Outcome Data.  _J. Amer. Statist. Assoc._ 110 (04 2015), 0–0.  https://doi.org/10.1080/01621459.2015.1023805

##  Appendix A Details on Data Simulation

###  A.1. Steps to Generate Semi-Synthetic Data

Algorithm 3 Simulate Semi-Synthetic Data

1:Randomly select 2N ($N=10Million$) users from the whole population and split
them equally into a training and testing dataset.

2:With the training dataset, train a regression model $\widehat{g_{j}}(X_{i})$
to predict the outcome metric $Y_{ij}$, a classification model
$\widehat{f}(X_{i})$ to predict the treatment assignment.

3:Simulate outcome metrics: generate predicted outcome
$\widehat{Y_{ij}}=\widehat{g_{j}}(X_{i})+\hat{\epsilon_{ij}}$, where
$\hat{\epsilon_{ij}}$ is sampled with replacement from the fitted errors of
the first $N$ used for training.

4:Simulate treatment assignment: Generate propensities
$Prob(T_{i}=1)=\widehat{f}(X_{i})$ for each user. Assign each of the last $N$
user into treatment or control by their predicted treatment status, which is
determined by a $Bernoulli(p)$ with $p=Prob(T_{i}=1)=\widehat{f}(X_{i})$.

###  A.2. Specification Details

####  A.2.1. Specification 1. Linear Model

  * •

Outcome models are fitted with linear regression, with 487 covariates.

  * •

Treatment selection model is fitted with binary logistic regression, with 487
covariates.

####  A.2.2. Specification 2. Linear Model with Interactions

We interact one user attribute, which is user persona, with all other
covariates. The new feature set contains 1458 covariates.

  * •

Outcome models are fitted with linear regression, with 1458 covariates.

  * •

Treatment selection model is fitted with binary logistic regression, with 1458
covariates.

####  A.2.3. Specification 3. Random Forest

  * •

Outcome models are fitted with RandomForestRegressor on 487 covariates. Number
of Trees is 10. Max Depth is 10.

  * •

Treatment selection model is fitted with RandomForestClassifier on 487
covariates. Number of Trees is 10. Max Depth is 15.

##  Appendix B Details of Balance and Weight Stability Metrics

The standardized mean difference (SMD):
$SMD_{d}=\frac{|\Delta_{d}|}{\sqrt{(s_{d1}^{2}+s_{d0}^{w2})/2}}$, where
$\Delta_{d}$ is the difference in means of covariate $d$ between treated group
and the re-weighted control group, $s_{d1}^{2}$ is the sample variance of
$X_{d}$ in the treated group, $s_{d0}^{w2}$ is the weighted sample variance of
$X_{d}$ in the re-weighted control group. To summarize the performance of all
dimensions, we further calculate $\overline{SMD_{first\>moment}}$,
$\overline{SMD_{square}}$, and $\overline{SMD_{interaction}}$ to represent the
mean $SMD$ of all covariates, all quadratic terms of the continuous variables,
and all possible interaction terms. Even though we do not add quadratic and
interaction terms in the weighting process, we are interested in the balance
performance of different methods in these terms. Most researchers consider
balance to be achieved when SMD is below 0.1. It is an arbitrary threshold.

The Mahalanobis balance (MB):
$MB=(\overline{X_{1}}-\overline{X_{0}^{w}})^{{}^{\prime}}\sum^{-1}(\overline{X_{1}}-\overline{X_{0}^{w}})$,
where $\overline{X_{1}}$ is the vector of covariate means in the treated
group, $\overline{X_{0}^{w}})$ is the vector of weighted covariate means in
the re-weighted control group, and $\sum^{-1}$ is the sample variance-
covariance matrix of covariates. A Mahalanobis Distance of 1 or lower shows
that the point is right among the benchmark
points.555https://www.theinformationlab.co.uk/2017/05/26/mahalanobis-distance/

The variance Ratio(VR): $VR_{d}=\frac{s_{d0}^{w2}}{s_{d1}^{2}}$. Closer values
to 1 indicate better balance in the second central moment. To make it
comparable with other balance metrics, we generally consider $|1-VR_{d}|$. To
summarize the performance of all dimensions, we calculate
$\overline{|1-VR_{d}|}$ which is the mean of $|1-VR_{d}|$ across all
dimensions. For good balance, the variance ratio should be between 0.5 and 2
(Rubin, 2001; Stuart, 2010). In our simulations, we use 0.5 for
$\overline{|1-VR_{d}|}$.

The Overlapping coefficient (OVL) is the proportion of overlap in the
covariate distributions between two groups. It is calculated by finding the
area under the minimum of both curves.

The OVL value is always between 0 and 1. Closer to 1 means better balance. To
make it comparable with other metrics, we focus on $1-OVL$. We have
$1-OVL_{d}$ for each dimension. To summarize the performance of all
dimensions, we calculate $\overline{1-OVL_{d}}$ which is the mean of
$1-OVL_{d}$ across all dimensions. There is no recommended threshold in
literature. We use 0.1 for $\overline{1-OVL_{d}}$.

The Kolmogorov-Smirnov distance (KS) is the maximum vertical distance between
two cumulative distributions. Some literature use 0.1 as the balance threshold
for KS (Markoulidakis et al., 2022).

The Standard Deviation of Normalized Weights (SD): $SD=sd(w_{i0})$ with
$\sum(w_{i0})=1$. It is the moment-based measures of variability. Lower value
indicates better weight stability.

The extremity of weights: measured by the maximum (MAX) and 99th percentile
(P99) of weights. Lower value indicates better weight stability.

The Effective Sample Size (ESS): $ESS=\frac{(\sum w_{i0})^{2}}{\sum
w_{i0}^{2}}$. Higher value indicates better weight stability.

##  Appendix C Further Details on Experiments

We provide critical code snippets for both the in-memory and Spark
implementation of EB and MicroSynth here:
https://gist.github.com/xizzzz/02caadbaf760510e267cc1fa9d3a7971

####  C.0.1. Learning Rate

We initialize learning rate with 0.01, and add decay if oscillation is
detected.

####  C.0.2. Stopping Condition

Heuristically we choose a threshold that makes our model achieve desirable
performance as the stopping condition. We used the difference between the sum
of the weights and the target below $1\text{\times}{10}^{-4}$.

####  C.0.3. Training Machine Types

We use Google Cloud Platform general purpose machines for the
experiments.666https://cloud.google.com/compute/docs/general-purpose-machines
Different machine types are chosen to fit the computation needs of each data
scale. Image version is 2.0.29-debian10. We list the machine types in the
following table. Data scale represents the number of rows in each data set.

Data Scale | Machine Type | Memory (GB)
---|---|---
10K | n1-standard-8 | 30
100K | n1-standard-16 | 60
1M | n2-highmem-32 | 256
10M | e2-highmem-8 | 64 * number of workers

##  Appendix D IPW hyperparameter tuning

In Section 4 we use grid search to tune hyperparameters of logistic regression
in the IPW method. We tune 3 hyperparameters: number of folds in cross-fitting
$k\in\\{0,3,5\\}$, regularization parameter
$C\in\\{0.0001,0.0003,0.001,0.003,0.01,0.03\\}$, and l1_ratio
$l\in\\{0,0.25,0.5,0.75,1\\}$. For each simulation setting, the best
hyperparameter is chosen with the lowest standardized mean difference.