Learning representations for counterfactual inference from observational data is of high practical relevance for many domains, such as healthcare, public policy and economics. Counterfactual inference enables one to answer "What if...?" questions, such as "What would be the outcome if we gave this patient treatment t 1 ?". However, current methods for training neural networks for counterfactual inference on observational data are either overly complex, limited to settings with only two available treatment options, or both. Here, we present Perfect Match (PM), a method for training neural networks for counterfactual inference that is easy to implement, compatible with any architecture, does not add computational complexity or hyperparameters, and extends to any number of treatments. PM is based on the idea of augmenting samples within a minibatch with their propensity-matched nearest neighbours. Our experiments demonstrate that PM outperforms a number of more complex state-of-the-art methods in inferring counterfactual outcomes across several real-world and semi-synthetic datasets.
========= Summary =========
The authors propose a novel method for counterfactual inference (i.e. individual/heterogeneous treatment effect, as well as average treatment effect) with neural networks. They perform propensity score matching within each minibatch in order to match the covariate distributions during training, which leads to a doubly robust model.
PM is evaluated on several standard semi-synthetic datasets (jobs, IHDP, TCGA) and PM shows state-of-the-art performance on some datasets, and overall looks quite promising. 
======= Comments =======
The paper is well-written, presents a novel method of some interest to the community, and shows quite good performance across a range of relevant benchmarks.
I have one major issue with this work: I don't see why propensity-score matching *within* a minibatch should provide a substantial improvement over propensity-score matching across the dataset (Ho et al 2011). I find the cursory explanation given ("it ensures that every gradient step is done in a way that is approximately unbiased") unconvincing, since (a) proper SGD training should be robust to per-batch biases during training (the expected loss is identical for both methods, correct?), and (b) biases should go away in the limit of large batch sizes. If indeed SGD required unbiased *minibatches* then standard minibatch SGD wouldn't work at all.
Looking at the experimental details in the appendix, it appears that the MatchIt package was used to do PSM, rather than a careful comparison under the same conditions. Are the exact matching procedure, PS estimator model, choosing "one of 6 closest  matches by propensity score", batch size, etc. the same between your PM implementation and MatchIt? I'd be very curious to see the results of a controlled comparison between Alg S1 and S2 under the same conditions (i.e. run your PM implementation on the whole dataset), and perhaps even some more clever experiments illustrating why matching within a minibatch is important. 
Another hypothesis for why PM is better than PSM is that the matching distribution for PM changes at each epoch (at least due to the randomization among the 6 closest matches). Could it be that the advantage of PM is that it actually provides a randomized rather than constant distribution of matched points?
Can the authors provide more motivation for why PM should outperform PSM? Or some more careful comparison of these methods isolating the benefits of PM? I think a convincing justification and comparison here could change my opinion, as I like the paper otherwise. Thanks!
Detailed Comments:
- There is insufficient explanation of the PM method in the main text. The method is only mentioned in a single sentence buried in the middle of a long paragraph "In PM, we match every sample within a minibatch...". This should be made more clear, e.g. by moving Algorithm S1 to the main text.
- The discussion on Model Selection and the argument for nearest-neighbor PEHE is clever and well-supported by the experiments.
- In Table 3 and 4, it's not clear which numbers are reported by the original authors and which were replicated by the authors.