summaries: Summary:  This paper considers settings where decisions for individuals are each made by different decision makers (agents). The paper defines the problem of finding a region in feature space in which the agents administer highly varying decisions. For example, perhaps judges mete out different sentences for misdemeanors by young men. The paper defines new causal criteria to characterize such regions in the feature space, where the choice of agent has an effect on the decision. The paper introduces an iterative algorithm to maximize this criteria to identify regions of heterogeneity, and provides theoretical results about the optimality of the proposed algorithm. The paper empirically demonstrates the algorithm on a semi-synthetic dataset and real clinical data about doctor recommendations for diabetic patients.  Contributions:  + A new criteria for defining regions where there is heterogeneity in decisions administered across agents. + Results for estimating the criteria from observational data and an iterative algorithm for optimizing the criteria. + Theoretical results about the algorithm and empirical validation of the algorithm.

summaries: In this paper, the authors develop an algorithm for identifying cases (e.g., defendants or patients) on which decision makers (e.g., judges and doctors) "disagree" on the decision.   ## Identification The authors consider a setting in which cases are summarized by features $X$. Each case is associated with a decision maker $A$ and a binary decision $Y$. In the pretrial release setting, each case is a defendant, the features $X$ are observable characteristics of the defendant (such as demographics or charge information), the decision maker $A$ is the assigned judge and the binary decision $Y$ is to either release or detain the defendant. The authors define the potential decision Y(a), which describes the decision that would be made by decision maker $a$.   The authors propose a summary measure of agent disagreement using average contrasts of potential decisions across agents given sets of the observable features. In particular, the authors focus on the "conditional relative agent bias" as  $$ \mathbb{E}[Y(a) - Y(\pi(x) | A = a, X \in S], $$ where $Y(\pi(x))$ denotes the expected decision under the assignment mechanism of decision makers to cases. The authors key identification assumption is that decision makers are as-if randomly assigned to cases conditional on the features (i.e., $A \perp Y(a) \mid X$). This measures the extent to which decision maker a's average decisions over $X \in S$ differs from the average decision that is made under the assignment mechanism. In other words, this measure of disagreement depends on the assignment mechanism of agents.  The conditional relative agent bias measures the extent to which a single agent a's decision disagrees relative to the average decision over $X \in S$. The authors then propose aggregating this measure across agents. For some grouping of agents G, the authors define the estimand $$ Q(S, G) = \sum_{a : G(a) = 1} P(A = a \mid X \in S) \mathbb{E}[Y(a) - Y(\pi(x) | A = a, X \in S]. $$ The authors main identification result (Theorem 1) is to show that if agents are as-if randomly assigned to cases, then this estimand can be identified as the conditional covariance between the observed decision and the agent assignment $$ Q(S, G) = \mathbb{E}[Cov(Y, G \mid X ) \mid X \in S]. $$  ## Estimating Regions of Heterogeneity. Using this identification result, the authors goal is to search over regions of the feature space $S$ and groupings over agents $G$ to find the combination that maximizes this proposed measure of disagreement. To do so, the authors propose to solve  $$ \max_{S, G} Q(S,G) \mbox{ s.t. } P(X \in S) \geq \beta, $$ which requires that the region have probability at least $\beta$ (a tuning parameter to be specified). The authors develop an iterative algorithm to optimize this objective function, and provide some conditions under which they are able to derive a generalization bound on the solution return after the first iteration (Assumptions 2-4 and Theorem 2).  ## Empirical Applications The authors apply their proposed algorithm in two empirical applications.   First, the authors analyze the behavior of their proposed algorithm in a semi-synthetic experiment that is calibrated to data from Lin et al. (2020) on recidivism risk predictions made by MTurk workers. The authors simulate data from two types of agents that follow different decision rules for detention vs. release. The authors compare their algorithm against direct models that fit logistic regressions for the decision on agent identifiers and observable features (direct models) and TARNet. The authors report that their algorithm does better at recovering the "region AUC."  Second, the authors analyze an observational dataset on medical treatment decisions for diabetes. The data contain information on 3,956 patients and 458 group practices, and the authors search for the region of disagreement across these group practices.

summaries: The authors are concerned with identifying regions in a covariate space. A region is characterised by having high variation in (potential) treatment outcome, as compared to the expected (potential) outcome. A method such as this could be used to identify variation of practice in the legal or medical domain (as indicated by the authors).   This is a relevant topic which the authors rightfully relate to the causal literature. Specifically, the authors chose the potential outcomes model, where the treatment is now a clinician or judge; the outcome is binary (amounting to a decision, e.g. bail or no bail); and the condition is an individual. Contrasting potential outcomes, the authors are interested in learning subsets of X, rather than Y(a).

summaries: This paper proposes a new method to identify "regions of disagreement" in a setting with human decision-makers assigned to individuals (eg. judges assigned to court cases). The "regions" are framed as subgroups of the data (eg. males over 30 years old) which can be interpretable (based on the chosen model). The paper formalizes this task as that of identifying the subgroup where the identity of the specific decision-maker has a causal effect on the decision, proposes an optimization algorithm to find this subgroup, develops generalization bounds (under some simplifying assumptions), and demonstrates compelling performance on semi-synthetic and real-world datasets.

abstractText: Individuals often make different decisions when faced with the same context, due to personal preferences and background. For instance, judges may vary in their leniency towards certain drug-related offenses, and doctors may vary in their preference for how to start treatment for certain types of patients. With these examples in mind, we present an algorithm for identifying types of contexts (e.g., types of cases or patients) with high inter-decision-maker disagreement. We formalize this as a causal inference problem, seeking a region where the assignment of decision-maker has a large causal effect on the decision. Our algorithm finds such a region by maximizing an empirical objective, and we give a generalization bound for its performance. In a semi-synthetic experiment, we show that our algorithm recovers the correct region of heterogeneity accurately compared to baselines. Finally, we apply our algorithm to real-world healthcare datasets, recovering variation that aligns with existing clinical knowledge.

authors:
Jason Abaluck
Leila Agha
Chris Kabrhel
Ali Raja
Arjun Venkatesh.

title: The determinants of productivity in medical testing: Intensity and allocation of care

venue: American Economic Review 106, 12 (2016), 3730–64.

authors:
Anna Aizer
Joseph J Doyle Jr.

title: Juvenile Incarceration, Human Capital and Future Crime: Evidence from Randomly-Assigned Judges

venue: Technical Report w19102. National Bureau of Economic Research.

title: Standards of medical care in diabetes—2010

venue: Diabetes care 33, Supplement 1 (2010), S11–S61.

authors:
David Arnold
Will Dobbie
Crystal S Yang.

title: Racial Bias in Bail Decisions

venue: The Quarterly Journal of Economics 133, 4 (May 2018), 1885–1932.

authors:
Femke Atsma
Glyn Elwyn
Gert Westert.

title: Understanding unwarranted variation in clinical practice: a focus on network effects, reflective medicine and learning health systems

venue: International Journal for Quality in Health Care 32, 4 (jun 2020), 271–274. https://doi.org/10.1093/ intqhc/mzaa023

authors:
John D Birkmeyer
Bradley N Reames
Peter McCulloch
Andrew J Carr
W Bruce Campbell
John E Wennberg.

title: Understanding of regional variation in the use of surgery

title: Pattern recognition and machine learning

venue: springer.

authors:
Michael D Cabana
Cynthia S Rand
Neil R Powe
Albert W Wu
Modena H Wilson
Paul-André C Abboud
Haya R Rubin.

title: Why don’t physicians follow clinical practice guidelines?: A framework for improvement

authors:
David C Chan Jr
Matthew Gentzkow
Chuan Yu.

title: Selection with variation in diagnostic skill: Evidence from radiologists

venue: Technical Report. National Bureau of Economic Research.

authors:
Ashley N Corallo
Ruth Croxford
David C Goodman
Elisabeth L Bryan
Divya Srivastava
Therese A Stukel.

title: A systematic review of medical practice variation in OECD countries

authors:
Janet Currie
W Bentley MacLeod.

title: Diagnosing expertise: Human capital, decision making, and performance among physicians

venue: Journal of labor economics 35, 1 (2017), 1–43.

authors:
Sanjoy Dasgupta
Leonard J Schulman.

title: A probabilistic analysis of EM for mixtures of separated, spherical Gaussians

venue: Journal of Machine Learning Research 8 (2007), 203–226.

authors:
Judith D De Jong
Gert P Westert
Ronald Lagoe
Peter P Groenewegen

title: Variation in hospital length of stay: do physicians adapt their length of stay decisions to what is usual in the hospital where they work

venue: Health Services Research 41,

authors:
Rafael Di Tella
Ernesto Schargrodsky.

title: Criminal Recidivism after Prison and Electronic Monitoring

venue: The Journal of Political Economy 121, 1 (2013), 28–73.

authors:
Julia Dressel
Hany Farid.

title: The accuracy, fairness, and limits of predicting recidivism

venue: Science advances 4, 1 (Jan. 2018), eaao5580.

authors:
Kenji Fukumizu
Arthur Gretton
Xiaohai Sun
Bernhard Schölkopf

title: Kernel measures of conditional dependence

venue: In NeurIPS,

authors:
Andrew Gelman
Jennifer Hill.

title: Data Analysis Using Regression and Multilevel/Hierarchical Models

venue: Cambridge University Press.

authors:
Christopher G Goetz
Stanley Fahn
Pablo Martinez-Martin
Werner Poewe
Cristina Sampaio
Glenn T Stebbins
Matthew B Stern
Barbara C Tilley
Richard Dodel
Bruno Dubois

title: Movement Disorder Society-sponsored revision of the Unified Parkinson’s Disease Rating Scale (MDS-UPDRS): process, format, and clinimetric testing plan

venue: Movement disorders 22,

authors:
Group
Henk Bilo
Luis Coentrão
Cécile Couchoud
Adrian Covic
Johan De Sutter
Christiane Drechsler
Luigi Gnudi
David Goldsmith
James Heaf

title: Clinical practice guideline on management of patients with diabetes and chronic kidney disease stage 3b or higher (eGFR< 45 mL/min). Nephrology Dialysis Transplantation

title: Long-term effectiveness of dopamine agonists and monoamine oxidase B inhibitors compared with levodopa as initial treatment for Parkinson’s disease (PD MED): a large, open-label, pragmatic randomised trial

authors:
M A Hernán
J M Robins.

title: Causal Inference: What If

venue: Chapman & Hall/CRC, Boca Raton.

authors:
George Hripcsak
Patrick B Ryan
Jon D Duke
Nigam H Shah
Rae Woong Park
Vojtech Huser
Marc A Suchard
Martijn J Schuemie
Frank J DeFalco
Adler Perotte

title: Characterizing treatment pathways at scale using the OHDSI network

venue: Proceedings of the National Academy of Sciences 113,

authors:
Jerry Kang
Judge Mark Bennett
Devon Carbado
Pam Casey
Nilanjana Dasgupta
David Faigman
Rachel Godsil
Anthony G. Greenwald
Justin Levinson
Jennifer Mnookin.

title: Implicit bias in the courtroom

authors:
David R Karger
Sewoong Oh
Devavrat Shah.

title: Budget-optimal task allocation for reliable crowdsourcing systems

venue: Operations Research 62, 1 (2014), 1–24.

authors:
Jon Kleinberg
Himabindu Lakkaraju
Jure Leskovec
Jens Ludwig
Sendhil Mullainathan.

title: Human Decisions and Machine Predictions

venue: The Quarterly Journal of Economics 133, 1 (Feb. 2018), 237–293.

title: Incarceration Length, Employment, and Earnings

venue: The American economic review 96, 3 (June 2006), 863–876.

authors:
Hoyt Koepke
Mikhail Bilenko.

title: Fast prediction of new feature utility

venue: arXiv preprint arXiv:1206.4680 (2012).

authors:
Zhiyuan (Jerry) Lin
Jongbin Jung
Sharad Goel
Jennifer Skeem

title: The Limits of Human Predictions of Recidivism

venue: Science Advances 6,

authors:
Kenneth Marek
Danna Jennings
Shirley Lasch
Andrew Siderowf
Caroline Tanner
Tanya Simuni
Chris Coffey
Karl Kieburtz
Emily Flagg
Sohini Chowdhury

title: The parkinson progression marker initiative (PPMI)

venue: Progress in neurobiology 95,

authors:
Sharon Muzerengi
Carl E Clarke.

title: Initial drug treatment in Parkinson’s disease

authors:
Xinkun Nie
Stefan Wager.

title: Quasi-oracle estimation of heterogeneous treatment effects

venue: arXiv preprint arXiv:1712.04912 (2017).

title: Examiner inconsistency: Evidence from refugee decisions

venue: Technical Report. Working paper

authors:
David L Ortega
H Holly Wang
Laping Wu
Nicole J Olynk.

title: Modeling heterogeneity in consumer preferences for select food safety attributes in China

venue: Food Policy 36, 2 (2011), 318–324. https://www.sciencedirect.com/science/article/pii/S0306919210001442

authors:
Adam Paszke
Sam Gross
Francisco Massa
Adam Lerer
James Bradbury
Gregory Chanan
Trevor Killeen
Zeming Lin
Natalia Gimelshein
Luca Antiga

title: Pytorch: An imperative style, high-performance deep learning

title: Causality: Models, Reasoning, and Inference (2nd ed.)

authors:
F. Pedregosa
G. Varoquaux
A. Gramfort
V. Michel
B. Thirion
O. Grisel
M. Blondel
P. Prettenhofer
R. Weiss
V. Dubourg
J. Vanderplas
A. Passos
D. Cournapeau
M. Brucher
M. Perrot
E. Duchesnay.

title: Scikit-learn: Machine Learning in Python

venue: Journal of Machine Learning Research 12 (2011), 2825–2830.

title: A scalable conditional independence test for nonlinear, non-gaussian data

venue: arXiv preprint arXiv:1401.5031 (2014).

title: EconML: A Python Package for ML-Based Heterogeneous Treatment Effects Estimation

authors:
Michael A Ribers
Hannes Ullrich.

title: Machine Predictions and Human Decisions with Variation in Payoffs and Skill

title: Conditional independence testing based on a nearest-neighbor estimator of conditional mutual information

venue: International Conference on Artificial Intelligence and Statistics. PMLR, 938–947.

authors:
Riccardo Scarpa
George Philippidis
Fiorenza Spalatro.

title: Product-country images and preference heterogeneity for Mediterranean food products: A discrete choice framework

authors:
Rajen D. Shah
Jonas Peters.

title: The Hardness of Conditional Independence Testing and the Generalised Covariance Measure

authors:
Shai Shalev-Shwartz
Shai Ben-David.

title: Understanding Machine Learning: From Theory to Algorithms

venue: Vol. 9781107057. Cambridge University Press. https://doi.org/10.1017/ CBO9781107298019

authors:
Uri Shalit
Fredrik D Johansson
David Sontag.

title: Estimating individual treatment effect: generalization bounds and algorithms

venue: International Conference on Machine Learning. PMLR, 3076–3085.

authors:
Le Song
Alex Smola
Arthur Gretton
Karsten M Borgwardt
Justin Bedo.

title: Supervised feature selection via dependence estimation

venue: Proceedings of the 24th international conference on Machine learning. 823–830.

authors:
AA Tahrani
GI Varughese
JH Scarpello
FWF Hanna

title: Metformin, heart failure, and lactic acidosis: is metformin absolutely contraindicated

authors:
Peter W Tuerk
Martina Mueller
Leonard E Egede.

title: Estimating physician effects on glycemic control in the treatment of diabetes: methods, effects sizes, and implications for treatment policy

authors:
Stefan Wager
Susan Athey.

title: Estimation and inference of heterogeneous treatment effects using random forests

title: seaborn: statistical data visualization

venue: Journal of Open Source Software 6, 60 (2021), 3021. https://doi.org/10.21105/joss.03021

title: All of Statistics: A Concise Course in Statistical Inference

authors:
Jacob Whitehill
Ting-fan Wu
Jacob Bergsma
Javier Movellan
Paul Ruvolo.

title: Whose vote should count more: Optimal integration of labels from labelers of unknown expertise

venue: Advances in neural information processing systems 22 (2009), 2035–2043.

authors:
Kun Zhang
Jonas Peters
Dominik Janzing
Bernhard Schölkopf.

title: Kernel-based conditional independence test and application in causal discovery

venue: arXiv preprint arXiv:1202.3775 (2012).

title: 2020) for the recidivism dataset

venue: We cite Marek et al

heading: 1 Introduction

text:
Understanding heterogeneity in decision-making is an established problem in medicine (Birkmeyer et al., 2013; Corallo et al., 2014; De Jong et al., 2006), consumer choice (Ortega et al., 2011; Scarpa et al., 2005), and law (Kang et al., 2012; Kleinberg et al., 2018; Arnold et al., 2018). In the context of medicine, this is referred to as the study of practice variation (Atsma et al., 2020; Cabana et al., 1999), where it is often observed that doctors, facing the same clinical context, will make different decisions. Likewise, in a legal context, judges often differ in their leniency in their decisions regarding bail (Kleinberg et al., 2018), juvenile incarceration (Aizer and Doyle Jr, 2013), the use of alternatives to incarceration (Di Tella and Schargrodsky, 2013), and incarceration length (Kling, 2006). In some scenarios this variation may be justified: The best medical treatment may not be obvious. In others, it may be grossly unfair, as in the case of racial bias in bail decisions (Arnold et al., 2018).
∗Equal contribution
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
In this work, we tackle the question of how to find and characterize this variation in the first place. In particular, we present a learning algorithm for identifying a “region of heterogeneity”, defined as a subset of all contexts (e.g., patients, cases) for which the identity of the decision-maker substantially affects the decision. In medicine, a better understanding of treatment variation can inform the development and dissemination of clinical guidelines. In the legal domain, characterizing the cases where judges vary most in their leniency may help with investigating potential issues of fairness.
We formalize characterizing the region of heterogeneity as a causal inference problem: We want to characterize examples where changing the decision-maker would have resulted in a different decision. The challenge is two-fold: First, we only observe a single decision-maker per example, so we cannot directly observe how (for instance) multiple judges would have decided the same case. Second, our data on individual decision-makers is often scarce. For instance, in Section 5, we consider a medical dataset with more than 400 doctors, each of whom has fewer than 9 samples on average.
We will refer to decision-makers as “agents”, and our contributions are as follows: In Section 2, we propose an objective defined in terms of counterfactual decisions across different agents, and show that this objective can be identified from observational data. Moreover, this objective does not require the use of agent-specific statistical models, making it amenable to our sparse setting. In Section 3, we give an iterative algorithm to identify regions of disagreement by maximizing this objective and provide intuition (in the form of a generalization bound) for the factors that drive its performance. In Section 4, we use a semi-synthetic dataset, derived from crowd-sourced recidivism predictions, to demonstrate that our algorithm recovers the correct region of heterogeneity accurately, even when there are many agents. Finally, in Section 5, we apply our algorithm to a real-world healthcare dataset and confirm that it recovers intuitive regions of variation in first-line diabetes treatment. We conclude with a discussion of related work and implications. Our code is available at https://github.com/clinicalml/finding-decision-heterogeneity-regions.
Our algorithm does not determine whether variation is inherently good or bad or how it should be addressed. Rather, more careful study with domain experts would be required to determine if variation can (or should be) reduced and how. In addition, false discovery of variation is possible and could have a negative impact. We expect that validation on independent datasets would be required in real-world applications, using the regions identified by our method as plausible hypotheses to test.

heading: 2 Characterizing Heterogeneity from a Causal Perspective

heading: 2.1 Notation

text: Let the data be drawn from a distribution P(X, A,Y), where X is a random variable representing context (or features), A is a discrete agent, and Y ∈ {0, 1} is the binary decision. The spaces of all X and A are denoted as X andA with realized values as lower case x and a, respectively. Indicator variables 1 [·] are one if the statement inside the brackets is true and zero otherwise. For a subset S ⊆ X, 1 [x ∈ S ] is sometimes written as a function S (x), where S : X → {0, 1}. A subset S may have several disjoint regions. E [Y |X ∈ S ] denotes the average Y across samples in S . For instance, if S = {X : X0 < 10}, then E [Y |X0 < 10] is a scalar average of Y among samples with X0 < 10.

heading: 2.2 Heterogeneity as a Causal Contrast

text:
Our conceptual goal is to identify a region S ⊆ X where different agents tend to make different decisions even when faced with the same context. We can formalize this in the language of potential outcomes from the causal inference literature (Pearl, 2009; Hernán and Robins, 2020), which for clarity we will refer to as potential decisions: In particular, we denote Y(a) to be the potential decision of agent a. The fundamental challenge of causal inference is that we do not observe all potential decisions {Y(a) : a ∈ A} for each sample, but only a single decision Y . With this in mind, we will make the following assumptions, standard in the literature on causal effect estimation.
Assumption 1 (Causal Identification Assumptions). (i) Consistency: Y = y, A = a =⇒ Y(a) = y, and (ii) No Unmeasured Confounding (NUC): For all a ∈ A, Y(a) ⊥ A | X.
Consistency links the potential Y(a) to the observed Y , and NUC says that there are no unobserved factors that influence both the assignment of agents and the decision itself. For instance, the quasirandom assignment of cases to judges conditioned on features X satisfies NUC (Kleinberg et al., 2018).
NUC may be violated if key aspects of the case (e.g., misdemeanor vs. felony) are omitted as features. For instance, misdemeanor and felony cases may be seen by different judges and have different decision processes, but this variation is not due to agent preferences. Given these assumptions, we propose a causal measure of agent-specific bias, defined as a contrast between potential decisions. Definition 1 (Conditional Relative Agent Bias). For an agent a ∈ A and a subset S ⊆ X, the conditional relative agent bias is defined as
E[Y(a) − Y(π(x)) | A = a, X ∈ S ] (1) where Y(a) is the potential decision of agent a, and Y(π(x)) B ∑ a′ E[Y(a′) | x]π(a′ | x) denotes the expected potential decision under the agent assignment distribution π(a′ | x) B P(A = a′ | X = x).
Note that under Assumption 1, Y(π(x)) = E[Y | x],2 but here we emphasize its causal interpretation as the expected decision of a random agent. Equation (1) represents the relative difference between the decision of an agent (on their particular distribution of cases in the region) and the potential decision of a random agent. In particular, Equation (1) can be written as follows under Assumption 1
E[Y(a) − Y(π) | A = a, X ∈ S ] = ∫
x∈S E[Y(a) − Y(π) | X = x]p(x | A = a, X ∈ S )dx, (2)
where we shorten Y(π(x)) to Y(π). This is the average difference (over p(x | a), restricted to those x in the set S ) of the conditional expected difference between Y(a) and Y(π). For example, suppose that the agent a is a judge who is particularly lenient on bail decisions for felony arrests (the region S ), and Y = 1 denotes granting bail. Then imagine the following counterfactual: Take the felony cases that are assigned to this judge and reassign each individual case, described by x, to a random judge a′, proportionally to p(a′ | x). We may then observe, on average, that the bail rate would decrease, because most judges are less lenient than judge a, corresponding to a positive value of Equation (1).
Equation (1) has the additional advantage of being easy to estimate: Under Assumption 1, it can be rewritten3 as E[Y−E[Y | X] | A = a, X ∈ S ], the expected residual in predicting (using the conditional expectation E[Y | X]) the decision of an agent a across the context x typically seen by that agent.

heading: 2.3 Formalizing a Causal Objective

text:
Our primary goal is to discover a region S where substantial heterogeneity exists across agents. To do so, we define an aggregate objective across a group G of agents, where G(a) ∈ {0, 1} is an indicator function for membership.
Q(S ,G) B ∑
a:G(a)=1
P(A = a | X ∈ S )E[Y(a) − Y(π) | A = a, X ∈ S ], (3)
We now show that this quantity can be identified and estimated from observational data without requiring agent-specific statistical models, before discussing the interpretation of this objective. Theorem 1 (Causal Identification). Under Assumption 1, Q(S ,G) can be identified as
Q(S ,G) = ES [Cov(Y,G | X)] = ES [(Y − E[Y | X])G], (4) where ES [·] B E[· | X ∈ S ] and Cov(Y,G | X) is the conditional covariance.
Theorem 1 and other theoretical results are proven in Appendix A. The result follows from proving that the agent-specific bias (Definition 1) is identifiable using the expected conditional covariance between Y and the binary indicator 1 [A = a]. With this in mind, we optimize the following objective, where the set S is constrained to be at least a certain size β and S is a hypothesis class of functions S .
max S∈S,G Q(S ,G) s.t., P(S ) ≥ β, (5)
Interpretation: Intuitively, this objective measures the disagreement between the agents in the group G(a) = 1 and the overall average E[Y | X] on the region S . Hence, the choice of group is important for interpreting the objective: If G(a) = 1 for all agents, the objective will be zero for any set S , as can be seen from Equation (4), applying the definition of the conditional expectation.
2See Proposition A1 for a short proof, and Proposition A2 for the derivation of Equation (2). 3See Proposition A3 in Appendix A.1.
Accordingly, we seek a region S for which the partially maximized objective L(S ) B maxG Q(S ,G) is large: This partial maximization is obtained by taking G(a) = 1 whenever the conditional relative agent bias of agent a (on the set S ) is non-negative. Thus, Equation (5) can be re-written as
max G Q(S ,G) = ∑ a∈A P(A = a | X ∈ S ) |E[Y(a) − Y(π) | A = a, X ∈ S ]|+ , (6)
where |x|+ B max(x, 0), and this objective becomes an average over agents who have a positive bias. This population objective is also equivalent (up to a constant factor) to the (weighted) sum of the magnitude of each agent’s conditional relative agent bias. See Proposition A4 in Appendix A.1.
Lack of Overlap: We have not made the overlap or positivity assumption that P(A = a | x) > 0 for all x, a. While this assumption is required to identify conditional causal effects E[Y(a) − Y(a′) | X] (Nie and Wager, 2017; Wager and Athey, 2018; Shalit et al., 2017), it is not required for identifying our causal contrast. Our problem only requires each context has a positive probability of being seen by more than one decision maker. For instance, suppose that S contains both misdemeanors and felonies and there are four judges a0, . . . , a3. If judges a0 and a1 make bail decisions exclusively for felonies while judges a2 and a3 make bail decisions exclusively for misdemeanors, our measure captures disagreement between a0 and a1 and between a2 and a3 even though comparisons between a0 and a2 or other pairs are impossible to make. Thus, we have chosen to compare Y(a) to the decisions of viable alternative agents, weighted by their probability p(a′ | x) of being selected for such a case.

heading: 3 Identifying Regions with Heterogeneity

text: In Section 3.1, we introduce an iterative optimization algorithm for a finite sample version of Objective (5) that alternatingly optimizes S and G. In Sections 3.2 and 3.3, we discuss practical heuristics for choosing the region size parameter β on training data and validating if the resulting region generalizes to held-out data. Finally, we build intuition for the factors that influence performance of this algorithm via a generalization bound in Section 3.4 under simplifying assumptions.

heading: 3.1 Iterative Optimization Algorithm

text:
We let Q̂(S ,G) be the empirical analog of Q(S ,G) (Equation 4), which we can write as follows,
Q̂(S ,G) B 1∑ a, j 1 [ xa j ∈ S ] ∑ a, j (ya j − f (xa j)) ·G(a) · 1 [ xa j ∈ S ] . (7)
where f (x) denotes a model of the conditional expectation f (x) ≈ E[Y | X = x]. For simplicity of notation, we assume that there are R samples (indexed by j) for each of a finite set of N agents (indexed by a), giving N · R samples in total. Our algorithm (Algorithm 1) takes as input the data {(xa j, ya j)} and a minimum region size β, and outputs a model h(x) and a threshold value b that describe a region of heterogeneity S = {x ∈ X; h(x) ≥ b}. Starting with S = X (the entire space), the algorithm identifies the grouping that maximizes Q̂(S ,G), then uses that grouping to identify the region maximizing the same quantity, repeating this process until convergence. The algorithm uses a classifier f (x) to estimate E[Y | X = x] and a regression model h(x) to estimate the conditional covariance of the decision Y and the grouping G at X = x. Note that we can use any supervised learning algorithms for f and h, allowing us to learn interpretable regions as part of the algorithm if h(x) is interpretable (e.g., decision trees). If sufficient data is available, samples can be split into three parts for estimating f (x) in line 2, computing G(a) in lines 5-8, and training h(x) and estimating the (1 − β)-th quantile in line 10. We do not perform this sample splitting because our sample sizes are already small. Under-fitting f (x) by further restricting the sample size could lead to false discovery if f (x) does not capture the variation explained by X.
Optimizing over G given S . Given a region S , our first result identifies the grouping G : A → {0, 1} that maximizes Q̂(S ,G) and shows that it can be expressed in terms of Q̂(S , 1 [A = a]). Proposition 1. Given S ⊆ X, Q̂(S ,G) is maximized over the space of functions G : A → {0, 1} at GS , where GS (a) = 1 [ Q̂(S , 1 [A = a]) ≥ 0 ] .
Intuitively, this proposition states that to maximize the empirical expected conditional covariance of the decision and grouping on a region, we must group agents by whether their residuals ya j − f (xa j) are (on average) positive or negative on S .
Algorithm 1 Identifying regions with variation
1: Input: Data {{xa j, ya j}Rj=1} N a=1 , minimum region size β. 2: Fit a model f (x) to E(Y | X = x). 3: Initialize S = X. 4: repeat 5: for a = 1, . . . ,N do 6: Compute Q̂(S , 1 [A = a]), where Q̂(S , 1 [A = a]) B 1∑ j 1[xa j∈S ] ∑ j(ya j − f (xa j))1 [ xa j ∈ S ] , 7: Set G(a) = 1 if Q̂(S , 1 [A = a]) ≥ 0 and 0 otherwise. 8: end for 9: Compute ba j = (ya j − f (xa j))G(a), a = 1, . . . ,N, j = 1, . . . ,R.
10: Fit a model h(x) to predict ba j from xa j, and let b be the (1 − β)-th quantile of h(xa j). 11: S ′ ← S . 12: S ← {xa j; h(xa j) ≥ b}. 13: until S = S ′ or iteration limit reached. 14: Output: Model h and threshold b, defining a region S = {x ∈ X; h(x) ≥ b}.
Optimizing over S given G. To optimize Q̂(S ,G) for a fixed grouping G over the hypothesis class S, we train a model h(x) to predict (ya j − f (xa j))G(a) given xa j, where h ∈ H . Using h as an estimate in Eq. 7, we find a set S to maximize the quantity 1∑ x 1[x∈S ] ∑
x∈S h(x), the empirical conditional expectation of h(x) over S . This quantity is maximized (subject to our β constraint) by including the largest β-fraction of the h(xa j) in S . Hence, we pick b as the (1 − β)-th quantile of h(xa j) and choose our region as Ŝ G = {x ∈ X; h(x) ≥ b}.

heading: 3.2 Tuning the Region Size Parameter

text: For real datasets, we need to choose β without knowledge of the true value. Given that our objective can be calculated on held-out data using the functions S ,G, a seemingly obvious approach would be to compute Q(S ,G) on a validation set and select the β that leads to the highest Q(S ,G). However, for a fixed data distribution, smaller values of β tend to produce higher values of Q(S ,G), and there is a trade-off between finding a smaller region of higher variation and a larger region that may include areas of lower (but still meaningful) variation. This motivated our original constraint P(S ) ≥ β. To select β, we propose a heuristic inspired by permutation-based hypothesis testing (Wasserman, 2004). We compare the training objective to a reference distribution of values (for the same β) that might be seen even if all agents followed the same policy. For each candidate β, we (i) run our algorithm and compute the objective on training data qobs B Q̂(Ŝ , Ĝ). (ii) For T iterations, we randomly shuffle the agents and re-run the algorithm to get a new objective value. This gives us a distribution P̂null over Q(S ,G) from a distribution where P(X,Y) and P(A) are unchanged but Y, X ⊥ A. (iii) Finally, we compute a p-value pβ = P̂null(Q > qobs) and choose the β with the smallest p-value. In Section 4.2, we find that this heuristic empirically recovers the true β value in semi-synthetic settings.

heading: 3.3 Validation of the Region

text:
We may wish to validate the learned region Ŝ independently of the grouping Ĝ. In particular, finding Ĝ is not our main goal, and we observe in our semi-synthetic experiments that our algorithm can find the true region S even when the grouping Ĝ is fairly poor (due to few samples per agent), as shown in Appendix B.2. We can optimize over G in Q(S ,G) to obtain an objective that depends only on S and can be used to compare regions. By Proposition 1, we obtain an empirical analog of Equation (6) as
L̂(S ) B max G
Q̂(S ,G) = 1∑ a, j 1 [ xa j ∈ S ] ∑ a ∣∣∣∣∣∣∣∣ ∑ j (ya j − f (xa j))1 [ xa j ∈ S ]∣∣∣∣∣∣∣∣ + , (8)
where |x|+ is equal to the positive part [x]+ = max(x, 0) as before. We then use this objective L̂(S ) to answer the following question: Does our chosen region Ŝ yield a significantly higher objective value on test data than a randomly selected region of the same size? An example of this analysis is given in Table 1 for the real-data experiment in Section 5.

heading: 3.4 Generalization Error

text:
We give a generalization bound for Algorithm 1 to build intuition for the factors that influence performance. To derive this bound, we consider a simplified setting, where there exists a set S ′,G′ such that the following set of assumptions hold. Assumption 2 (Group-based variation). For all x ∈ S ′, E[Y | X = x, A = a] = E[Y | X = x,G′(a)] and for all x < S ′, E[Y | X = x, A = a] = E[Y | X = x] Assumption 3 (Non-zero relative biases). There exists a constant α > 0 such that for all x ∈ S ′, E[Y | X = x,G′(A) = 1] − E[Y | X = x] > α, and E[Y | X = x,G′(A) = 0] − E[Y | X = x] < −α, Assumption 4 (All agents see samples in S ′). There exists a constant ω > 0, such that for every a ∈ A, P(X ∈ S ′ | A = a) > ωP(X ∈ S ′).
Note that under these assumptions, S ′,G′ maximize the objective Q(S ,G) (see Appendix A.4), so we will refer to them as S ∗,G∗ for the remainder of this section. Assumption 2 says that there are two groups of agents, who follow two distinct decision policies within a region S ∗ but follow an identical decision policy outside of S ∗. Assumption 3 says that one group has a positive bias across all of S ∗, relative to the average over both groups, and the other group has a negative bias. To simplify the analysis, we also make Assumption 4 that every agent has some non-zero chance of observing some contexts X in the region, but note that we do not require that p(x | a) > 0 for all x ∈ S ′. Under these assumptions, we demonstrate that the first iteration of Algorithm 1 will find, with high probability, a region Ŝ whose value Q(Ŝ ,G∗) (for the same grouping G∗ defined above) is close to that of the optimal S ∗. Note that we do not claim that the iterative algorithm finds the globally optimal solution. For simplicity, we assume that f (x) perfectly recovers E[Y | X]. This can be relaxed at the cost of additional terms in the bound that go to zero as the overall sample size increases. Under Assumptions 2, 3, and 4, assume that P(S ∗) = β. For the informal version presented here, we assume that P(S ∗) = P(Ŝ ) = β, where Ŝ is returned by our algorithm, and that exactly a β-fraction of our samples fall into S ∗ and Ŝ (in the Appendix, we give a version without these simplifications). Theorem 2. Under the assumptions above, if S ∗ ∈ S and R > 2 ln 2
α2β2ω2 , the first iteration of Algorithm 1
returns Ŝ such that, with probability at least 1 − δ, Q(S ∗,G∗) − Q(Ŝ ,G∗) ≤ , where
= √ 2 ln(3/δ) βN · R + 2 β η + √3η(1 − η)δ · N  + 1β 2R(S,N · R) + 4 √2 ln(12/δ)N · R  ,
where R(S,N · R) is the Rademacher complexity of S, and η B exp ( −Rα2β2ω2/2 ) .
The term η plays an important role: It bounds the expected misclassification error P(Ĝ(a) , G∗(a)). For sufficiently large R, we have that η < 1/2 with high probability, i.e., we have a better-than-random chance of identifying the group for an individual agent. Moreover, η decreases as we increase the number of samples R for each agent, the separation α between the two groups on S ∗, the region size β, and the constant ω. For sufficiently small η, our algorithm discovers a region whose value (in terms of Q(S ,G∗)) is close to that of the optimal region. The generalization bound improves as the number of agents N increases, the number of samples R for each agent increases, or the complexity of the hypothesis class decreases. The latter is measured here by the Rademacher complexity R(S,N · R) of our hypothesis class S, which can be bounded by standard arguments. In conclusion, under some additional assumptions, Algorithm 1 identifies an approximately optimal solution with high probability after one iteration. We show via semi-synthetic experiments in Appendix B.3 that convergence is generally fast in practice.

heading: 4 Semi-Synthetic Experiment: Recidivism Prediction

text:
For conceptual motivation in the introduction, we discussed the legal system: As a potential application of our method, one could determine types of cases for which the idiosyncratic preferences of judges have a significant impact on their decisions. Lacking data on judge decisions with sufficient context, we turn to the more controlled setting of human predictions of recidivism.
Dataset: We use publicly available data from Lin et al. (2020),4 who ask participants on Amazon’s Mechanical Turk platform to make recidivism predictions based on information present in the
4Available at https://github.com/stanford-policylab/recidivism-predictions
“Correctional Offender Management Profiling for Alternative Sanctions” (COMPAS) dataset for Broward County, FL (Dressel and Farid, 2018). Participants (or “agents”) are shown 5 risk factors: age, gender, number of prior convictions, number of juvenile felony charges, and number of juvenile misdemeanor charges. The charge in question is also given, as well as whether the charge is a misdemeanor or felony. The dataset contains 4550 cases evaluated by 87 participants.
Semi-Synthetic Policy Generation: To benchmark our method, we generate semi-synthetic data where we have access to a “ground truth” region of heterogeneity. We retain the features presented to the original participants and construct two policies, which we refer to as the “base” and “alternative” policies: For the base policy, we learn a logistic regression model on the binary decisions across the whole dataset. For the alternative policy, an extra positive term is added to the logistic regression for samples within the region. We construct two scenarios with different regions of variation: (1) all drug possession charges, and (2) all misdemeanor charges where the individual is 35 years old or younger. These make up 22% and 21% of the data, respectively. Then, we generate synthetic agents (randomly assigned to cases) and assign half of the agents to the base policy and half to the alternative. Synthetic decisions are then sampled from the logistic regressions. For each scenario, the two groups of agents follow the same stochastic policy outside of the region, and one group systematically prefers Y = 1 within the region. More details can be found in Appendix B.1.

heading: 4.1 Performance versus Baselines

text:
Baselines: We compare how well our approach identifies the true region of heterogeneity with several baselines. To our knowledge, the problem of finding regions of heterogeneity with a large number of agents has not been studied before. Many causal inference methods for treatment effect estimation are designed for a single, binary treatment. However, naively estimating the treatment effect between each pair of providers would scale O ( |A|2 ) . Therefore, we develop new baselines. Some (including
the causal forest and U-learner adaptations described in Appendix B.2) are based on causal inference methods augmented to identify a region of heterogeneity and grouping of agents where possible.
Direct models: This baseline measures how much adding the agent to the feature set improves prediction of decisions. We fit logistic regressions with and without the agent feature to estimate E[Y | A, X] and E[Y | X]. For each sample (x, y, a), we compute |y − E[Y | X = x]| − |y − E[Y | A = a, X = x]| to quantify how much the model with agents outperforms the model without agents. Then, we fit a “region model” to predict this quantity from X. This region model is either a ridge regression, decision tree, or random forest model. Finally, we compute the top β quantile of predictions from the region model in the training and validation sets and use this cut-off to select points in the test set.
TARNet: A treatment-agnostic representation network (Shalit et al., 2017) models the outcomes of all treatments for each sample by learning a shared representation of the sample features and then
having a separate prediction head for each treatment. We implement the shared representation model using a 2-layer neural network with ReLU and dropout. Each prediction head is a linear layer with a sigmoid function. TARNet predicts E[Y | X, A] for every agent for each sample. VarA[E[Y | X, A]] measures the variation across all agents if they had seen context X and is analogous to our objective without grouping. We predict this quantity with the region models as in the direct model baseline.
Results: We evaluate how well the algorithms identify the samples within the region of heterogeneity when we vary the number of agents among 2, 5, 10, 20, 40, and 87, where 87 is the number of agents in the original dataset. Figure 1 shows the best overall region models for each set-up, with the other models deferred to Appendix B.2. The metric in Figure 1 is the region AUC, defined as how well the model classifies whether samples belong in the region of heterogeneity when compared to the true region. Algorithm 1 consistently performs well for both semi-synthetic set-ups, especially when the number of agents increases to a realistic level (and the number of samples per agent decreases). The direct baseline deteriorates very rapidly as the number of agents increases in the drug possession set-up, while the TARNet baseline deteriorates rapidly in the misdemeanor under age 35 set-up. Refer to Appendix B.2 for additional baseline details, region models, and evaluation metrics. We also show that our method is robust in a set-up with more than 2 agent groups in Appendix B.4.

heading: 4.2 Tuning the Region Size

text: We validate the proposed approach of tuning β (discussed in Section 3.2), by applying the methodology to our semi-synthetic setting here. We sample 30 semi-synthetic datasets and consider candidate values of β in [0.02, 0.42] in increments of 0.04. For each proposed value of β we use T = 40 random permutations of the agents. Figure 2 presents results for the misdemeanor under age 35 set-up. As the candidate value of β increases (up to the true value of β), the p-value decreases, and the distribution of selected β values are centered on the true value of β.

heading: 5 Real-data Experiment: First-Line Diabetes Treatment

text:
We apply our algorithm on a real-world dataset consisting of first-line (initial) treatment for type 2 diabetes and examine how the treatment variation we discover aligns with clinical knowledge. We present an additional real-world experiment (on Parkinson’s disease) in Appendix D.
Data and Setup: We use an observational, de-identified, dataset provided by a large health insurer in the United States, spanning from 2010 to 2020. The task is to classify first-line treatment decisions between metformin (Y = 0)–the typical recommendation from the American Diabetes Association– and other common first-line treatments such as sitagliptin, glipizide, glimepiride, or glyburide (Y = 1) (American Diabetes Association, 2010; Hripcsak et al., 2016). As relevant clinical features, we include the patients’ most recent eGFR (mL/min/1.73m2) and creatinine (mg/dL) measurements,
incidence of heart failure, and treatment date. Because treatment date does not define a type of patient, we omit it from the region model. However, including it in the outcome model is essential because of increasing use of metformin over time. The agent indicator A is the group practice of the doctor responsible for the patient’s treatment. 3,980 patients and 447 group practices are included in our cohort. After requiring at least 4 patients per agent, 3,576 patients and 176 group practices are included. This filter ensures each group practice has at least 1 sample in the training and validation sets and at least 2 samples in the test set. In this experiment, we choose β = 0.25 as input to our algorithm. See Appendix C for additional cohort definition and set-up details.
Interpretation of Results: To interpret the region, we use decision trees as our region model h(x). The tree is visualized in Appendix C. The decision tree identifies the region of heterogeneity as the union of (i) eGFR below 71.5 and (ii) eGFR above 98.5 and creatinine above 0.815. These regions align with clinical intuition. In the first region, low eGFR values indicate impaired renal function (Group et al., 2009), which is a contraindication for metformin since it is traditionally associated with increased risk of lactic acidosis (Tahrani et al., 2007). Still, treatment decisions can vary here because guidelines for managing patients with eGFR below 45 are lacking (Group et al., 2015). Note that this region provides an example of how our algorithm works when overlap is not satisfied. Although 33 of 176 group practices do not see patients with these features, we can still conclude that this is a region of heterogeneity among the 143 agents with cases. In the second region, there are no obvious contraindications for metformin. Thus, understanding why some doctors on average only prescribe metformin 78% of the time to patients in this region may help us identify whether this is a region in which we can standardize practice.
Assessing Significance: In Table 1, we perform a sanity check, assessing whether our algorithm identifies a region S whose variation in held-out data is higher than that of a randomly selected region, using the partially optimized objective L(S ) = maxG Q(S ,G) laid out in Section 3.3 to compare regions directly. Table 1 shows that L(Ŝ ) is similar on the training and test data. Furthermore, we compare L(Ŝ ) on the test set to the distribution of L(S rand), where S rand are random subsets of the test data of the same size as Ŝ . We compute the latter distribution using 100 random subsets and observe that the test value of L(Ŝ ) is more than two standard deviations above the mean of the latter. This gives us confidence that the discovered region S generalizes to a region of heterogeneity beyond the training set. We direct the reader to Appendix C for additional analyses, such as evaluating the stability of the region over multiple folds of splitting the data.

heading: 6 Related Work

text:
Beyond previously mentioned connections to causal effect estimation, we highlight a few areas of research that share similar goals to our own.
Agent-specific Models of Decisions: Prior works have modeled agent-specific decision-making processes by estimating a separate model for each agent. Abaluck et al. (2016) model heterogeneity in physician tendency to run diagnostic tests. Chan Jr et al. (2019) estimate radiologist skill based on diagnosis and miss rates. In our setting, unlike in diagnosis, there is no “correct” decision that can be incorporated into the model. Ribers and Ullrich (2020) assume there is provider-specific noise in determining patient type, which affects the pay-off functions for deciding whether to prescribe antibiotics. When only a few decisions are observed per agent, these agent-specific models cannot be estimated reliably. Currie and MacLeod (2017) also incorporate physician beliefs and skill into a pay-off function. They estimate an aggregate logistic choice model (for C-sections) across all physicians and then learn how individual physicians deviate from this model. They do not learn the regions where this deviation occurs, as they focus on how heterogeneity is associated with
downstream outcomes. Norris (2020) looks for disagreement between agents but relies on some cases being seen by multiple agents. We assume each case is seen by only one decision-maker.
Conditional Independence Testing: While our objective maximizes a causal notion of dependence, one could instead ask if Y is conditionally independent of A given X. Many metrics exist for testing conditional independence, such as the Hilbert-Schmidt independence criterion (HSIC) (Fukumizu et al., 2007; Zhang et al., 2012), conditional mutual information (Runge, 2018), conditional correlation (Ramsey, 2014), and expected conditional covariance (Shah and Peters, 2020). We give the last a causal interpretation under some assumptions and seek a region that maximizes it, in lieu of testing.
Hierarchical /Mixture Models: Bayesian methods are often used to estimate group-level effects, such as a per-physician effect on patient outcomes (Tuerk et al., 2008), where group identifiers are included as a categorical feature in a multi-level generalized linear model (Gelman and Hill, 2006). Alternatively, one could assume a conditional mixture model (Bishop, 2006), where agents belong to latent clusters that each have their own policy. However, both of these methods require parametric assumptions on the distribution of P[Y | x, a], and even so, the optimal mixture model is not necessarily identifiable when both the clusters and policies are unknown (Dasgupta and Schulman, 2007). By contrast, our method does not require making particular parametric assumptions about E[Y | x, a] and seeks to learn the region of heterogeneity directly. Feature Evaluation: Checking for heterogeneity can also be framed as feature evaluation, where we would like to evaluate whether adding the agent identifier will increase predictive power. Feature evaluation methods typically maximize dependence between selected features and labels, utilizing measures similar to those in conditional independence testing, such as the HSIC (Song et al., 2007). Other methods use the correlation of the new feature with the loss gradient of the current predictor as a measure of utility (Koepke and Bilenko, 2012). In contrast, we focus not on checking marginally for the predictive power of agent identifiers, but rather identifying a region.
Crowdsourcing: Our work is conceptually related to identifying which samples are difficult to label via crowdsourcing annotations (Karger et al., 2014; Whitehill et al., 2009). Most crowdsourcing models are generative with latent variables for the correct sample labels, sample difficulty, agent expertise, and agent bias. They then optimize for the likelihood of the observed labels. The set of difficult samples is analogous to our region of heterogeneity. The main difference with our problem is that we do not require any notion of the “correct” label.

heading: 7 Discussion

text:
In this work, we take a causal perspective on finding regions where agents (i.e., decision-makers) have heterogeneous preferences, formalizing this heterogeneity in terms of counterfactual (or “potential”) decisions across agents. We propose a causal measure of agent-specific biases and give an objective that aggregates this bias over agents. This objective can be identified from observational data and written in terms of an expected conditional covariance. Importantly, for our applications, this does not require building agent-specific models or assuming overlap across all agents.
We give an iterative algorithm to find a region that maximizes this objective. Then, we demonstrate in semi-synthetic experiments that our algorithm accurately recovers the region of heterogeneity and scales well with the number of agents. In contrast, performance of baslines deteriorates when the number of agents increases. Although our experiments have low-dimensional spaces, we hypothesize our algorithm would scale well to high-dimensional feature spaces since the average policy and region models can handle high-dimensional input spaces. Finally, on a real-world medical dataset, we show that our algorithm can yield insights that align with clinical knowledge.
Our work is motivated by understanding variation in human decision-making. In the judicial domain, our method may help unearth types of cases where decisions are highly dependent on the judge assigned to the case. In the medical domain, our approach may identify types of patients where new guidelines may be required to help doctors make decisions more consistent with standard of care. Domain expertise is required to determine the implications of the regions discovered by our method. Beyond these applications, we see our approach as a useful data science tool for understanding heterogeneity in decisions that appears to be driven by individual-level preferences.
Acknowledgements: We would like to thank Aaron Smith-McLallen, James Denyer, Luogang Wei, Johnathon (Kyle) Armstrong, Neil Dixit, Aditya Sai, and the rest of the data science group at Independence Blue Cross, whose expertise, data, and support enabled the diabetes experiment. We would also like to thank Rebecca Boiarsky for converting the diabetes data to the OMOP common data model format, Monica Agrawal for her helpful comments, and other members of the lab for insightful discussions. We are also grateful to Charles Venuto and Monica Javidnia for their advice on Parkinson’s. This work was supported in part by Independence Blue Cross, Office of Naval Research Award No. N00014-17-1-2791, an Abdul Latif Jameel fellowship, and a NSF CAREER award.

title: Finding Regions of Heterogeneity in Decision-Making via Expected Conditional Covariance


summaries: This paper proposes TokenGAN for unconditional image synthesis. The token-based generator takes the learned content tokens and style tokens from the latent space as input, and output a series of style-modulated content tokens, which are then concatenated and reshaped to get the output image. Cross-attention is used to assign the styles to different content tokens. The model adopts the style-based generator idea from StyleGAN and introduces the recent token-based structures and transformers to the traditional style-based generator. Experiments are conducted on FFHQ and LSUN CHURCH datasets and the quantitative results are slightly better than StyleGAN2.

summaries: In this paper, a transformer-based modification of the generator in the ubiquitous StyleGAN model is proposed (named _TokenGAN_). More specifically, the entire generator is replaced by a convolution-free architecture where style-content modulation is now modeled using cross-attention layers instead of AdaIN. Similar to the original StyleGAN, this should allow for decoupled control of style and content at different granularities. The paper shows that the proposed approach achieves FID scores comparable to or better than StyleGAN2 for unconditional image synthesis on FFHQ and LSUN-Churches.

summaries: This paper proposes to adopt the StyleGAN architecture (i.e. individual styles that affect image generation) to a transformer based model for image generation. The model generates individual image patches which are modulated by style tokens through an attention mechanism. The model is evaluated on the FFHQ and LSUN Church datasets and performs similarly well or better than StyleGAN2 based on FID, precision, and recall.

summaries: This paper investigates the application of transformer for image generation by replacing the convolutional layer with MLP(including attention). The motivation provided is the limitation of the style-based generator (e.g.,  StyleGAN), which leverages the single style over multiple image regions, allowing a less entangled representation learning in W and visible artifacts in generated images.  To combat this shortcoming, authors introduce transformer-liked architecture, and successfully generate high-quality images (1024*1024) without touching any convolutional layers in generator.

abstractText: We present a new perspective of achieving image synthesis by viewing this task as a visual token generation problem. Different from existing paradigms that directly synthesize a full image from a single input (e.g., a latent code), the new formulation enables a flexible local manipulation for different image regions, which makes it possible to learn content-aware and fine-grained style control for image synthesis. Specifically, it takes as input a sequence of latent tokens to predict the visual tokens for synthesizing an image. Under this perspective, we propose a token-based generator (i.e.,TokenGAN). Particularly, the TokenGAN inputs two semantically different visual tokens, i.e., the learned constant content tokens and the style tokens from the latent space. Given a sequence of style tokens, the TokenGAN is able to control the image synthesis by assigning the styles to the content tokens by attention mechanism with a Transformer. We conduct extensive experiments and show that the proposed TokenGAN has achieved state-of-the-art results on several widelyused image synthesis benchmarks, including FFHQ and LSUN CHURCH with different resolutions. In particular, the generator is able to synthesize high-fidelity images with 1024× 1024 size, dispensing with convolutions entirely.

authors:
Rameen Abdal
Yipeng Qin
Peter Wonka

title: Image2stylegan: How to embed images into the stylegan latent space

venue: In ICCV,

title: Image2stylegan++: How to edit the embedded images

venue: In CVPR,

title: Adam: A method for stochastic optimization

venue: In ICLR,

authors:
Andrew Brock
Jeff Donahue
Karen Simonyan

title: Large scale gan training for high fidelity natural image synthesis

authors:
Tom B Brown
Benjamin Mann
Nick Ryder
Melanie Subbiah
Jared Kaplan
Prafulla Dhariwal
Arvind Neelakantan
Pranav Shyam
Girish Sastry
Amanda Askell

title: Language models are few-shot learners

authors:
Nicolas Carion
Francisco Massa
Gabriel Synnaeve
Nicolas Usunier
Alexander Kirillov
Sergey Zagoruyko

title: End-to-end object detection with transformers

venue: In ECCV,

authors:
Hanting Chen
Yunhe Wang
Tianyu Guo
Chang Xu
Yiping Deng
Zhenhua Liu
Siwei Ma
Chunjing Xu
Chao Xu
Wen Gao

title: Pre-trained image processing transformer

authors:
Mark Chen
Alec Radford
Rewon Child
Jeffrey Wu
Heewoo Jun
David Luan
Ilya Sutskever

title: Generative pretraining from pixels

venue: In ICML,

authors:
Edo Collins
Raja Bala
Bob Price
Sabine Susstrunk

title: Editing in style: Uncovering the local semantics of gans

authors:
Emily Denton
Soumith Chintala
Arthur Szlam
Robert Fergus

title: Deep generative image models using a laplacian pyramid of adversarial networks

venue: In NeurIPS,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: Bert: Pre-training of deep bidirectional transformers for language understanding

venue: In NAACL-HLT,

authors:
Chris Donahue
Zachary C Lipton
Akshay Balsubramani
Julian McAuley

title: Semantically decomposing the latent spaces of generative adversarial networks

authors:
Alexey Dosovitskiy
Lucas Beyer
Alexander Kolesnikov
Dirk Weissenborn
Xiaohua Zhai
Thomas Unterthiner
Mostafa Dehghani
Matthias Minderer
Georg Heigold
Sylvain Gelly

title: An image is worth 16x16 words: Transformers for image recognition at scale

authors:
Patrick Esser
Robin Rombach
Bjorn Ommer

title: Taming transformers for high-resolution image synthesis

authors:
Rinon Gal
Dana Cohen
Amit Bermano
Daniel Cohen-Or

title: Swagan: A style-based wavelet-driven generative model

venue: arXiv preprint arXiv:2102.06108,

authors:
Ian J Goodfellow
Jean Pouget-Abadie
Mehdi Mirza
Bing Xu
David Warde-Farley
Sherjil Ozair
Aaron C Courville
Yoshua Bengio

title: Generative adversarial nets

venue: NeurIPS,

authors:
Jinjin Gu
Yujun Shen
Bolei Zhou

title: Image processing using multi-code gan prior

authors:
Martin Heusel
Hubert Ramsauer
Thomas Unterthiner
Bernhard Nessler
Sepp Hochreiter

title: Gans trained by a two time-scale update rule converge to a local nash equilibrium

authors:
Zhiyuan Hu
Jia Jia
Bei Liu
Yaohua Bu
Jianlong Fu

title: Aesthetic-aware image style transfer

venue: In ACM MM,

authors:
Xun Huang
Serge Belongie

title: Arbitrary style transfer in real-time with adaptive instance normalization

authors:
Drew A Hudson
C. Lawrence Zitnick

title: Generative adversarial transformers

authors:
Ali Jahanian
Lucy Chai
Phillip Isola

title: On the "steerability" of generative adversarial networks

authors:
Yifan Jiang
Shiyu Chang
Zhangyang Wang

title: Transgan: Two pure transformers can make one strong gan, and that can scale up

authors:
Tero Karras
Timo Aila
Samuli Laine
Jaakko Lehtinen

title: Progressive growing of gans for improved quality, stability, and variation

authors:
Tero Karras
Miika Aittala
Janne Hellsten
Samuli Laine
Jaakko Lehtinen
Timo Aila

title: Training generative adversarial networks with limited data

authors:
Tero Karras
Samuli Laine
Timo Aila

title: A style-based generator architecture for generative adversarial networks

authors:
Tero Karras
Samuli Laine
Miika Aittala
Janne Hellsten
Jaakko Lehtinen
Timo Aila

title: Analyzing and improving the image quality of stylegan

authors:
Tuomas Kynkäänniemi
Tero Karras
Samuli Laine
Jaakko Lehtinen
Timo Aila

title: Improved precision and recall metric for assessing generative models

authors:
Ze Liu
Yutong Lin
Yue Cao
Han Hu
Yixuan Wei
Zheng Zhang
Stephen Lin
Baining Guo

title: Swin transformer: Hierarchical vision transformer using shifted windows

authors:
Sachit Menon
Alexandru Damian
Shijia Hu
Nikhil Ravi
Cynthia Rudin

title: Pulse: Selfsupervised photo upsampling via latent space exploration of generative models

authors:
Lars Mescheder
Andreas Geiger
Sebastian Nowozin

title: Which training methods for gans do actually converge

authors:
Takeru Miyato
Toshiki Kataoka
Masanori Koyama
Yuichi Yoshida

title: Spectral normalization for generative adversarial networks

authors:
Aäron van den Oord
Nal Kalchbrenner
Oriol Vinyals
Lasse Espeholt
Alex Graves
Koray Kavukcuoglu

title: Conditional image generation with pixelcnn decoders

authors:
Taesung Park
Ming-Yu Liu
Ting-Chun Wang
Jun-Yan Zhu

title: Semantic image synthesis with spatially-adaptive normalization

authors:
Stanislav Pidhorskyi
Donald A Adjeroh
Gianfranco Doretto

title: Adversarial latent autoencoders

authors:
Elad Richardson
Yuval Alaluf
Or Patashnik
Yotam Nitzan
Yaniv Azar
Stav Shapiro
Daniel Cohen-Or

title: Encoding in style: a stylegan encoder for image-to-image translation

authors:
Mehdi SM Sajjadi
Olivier Bachem
Mario Lucic
Olivier Bousquet
Sylvain Gelly

title: Assessing generative models via precision and recall

authors:
Yujun Shen
Jinjin Gu
Xiaoou Tang
Bolei Zhou

title: Interpreting the latent space of gans for semantic face editing

authors:
Dmitry Ulyanov
Andrea Vedaldi
Victor Lempitsky

title: Improved texture networks: Maximizing quality and diversity in feed-forward stylization and texture synthesis

authors:
Aaron Van Oord
Nal Kalchbrenner
Koray Kavukcuoglu

title: Pixel recurrent neural networks

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

authors:
Yi Wang
Ying-Cong Chen
Xiangyu Zhang
Jian Sun
Jiaya Jia

title: Attentive normalization for conditional image generation

authors:
Hongwei Xue
Bei Liu
Huan Yang
Jianlong Fu
Houqiang Li
Jiebo Luo

title: Learning fine-grained motion embedding for landscape animation

authors:
Fuzhi Yang
Huan Yang
Jianlong Fu
Hongtao Lu
Baining Guo

title: Learning texture transformer network for image super-resolution

authors:
Fisher Yu
Ari Seff
Yinda Zhang
Shuran Song
Thomas Funkhouser
Jianxiong Xiao

title: Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop

venue: arXiv preprint arXiv:1506.03365,

authors:
Yanhong Zeng
Jianlong Fu
Hongyang Chao

title: Learning joint spatial-temporal transformations for video inpainting

authors:
Yanhong Zeng
Jianlong Fu
Hongyang Chao
Baining Guo

title: Learning pyramid-context encoder network for high-quality image inpainting

authors:
Han Zhang
Ian Goodfellow
Dimitris Metaxas
Augustus Odena

title: Self-attention generative adversarial networks

authors:
Richard Zhang
Phillip Isola
Alexei A Efros
Eli Shechtman
Oliver Wang

title: The unreasonable effectiveness of deep features as a perceptual metric

authors:
Heliang Zheng
Jianlong Fu
Yanhong Zeng
Jiebo Luo
Zheng-Jun Zha

title: Learning semanticaware normalization for generative adversarial networks

authors:
Jiapeng Zhu
Yujun Shen
Deli Zhao
Bolei Zhou

title: In-domain gan inversion for real image editing

authors:
Peihao Zhu
Rameen Abdal
Yipeng Qin
Peter Wonka. Sean

title: Image synthesis with semantic region-adaptive normalization

authors:
Xizhou Zhu
Weijie Su
Lewei Lu
Bin Li
Xiaogang Wang
Jifeng Dai

title: Deformable detr: Deformable transformers for end-to-end object detection

heading: 1 Introduction

text:
Unconditional image synthesis generates images from latent codes by adversarial training [10, 16, 19, 25, 27, 33, 50]. Recent advances have been achieved by a style-based generator architecture in terms of both the visual quality and resolution of generated images [26, 27, 28, 37, 52]. In particular, the style-based generator has been widely used in many other generative tasks, including facial editing [9, 40], style transfer [1, 38], image super-resolution [17, 31], and image inpainting [2, 52].
The key to the success of the style-based generator lies in the learning of the style control based on the intermediate latent spaceW [27, 28]. Instead of feeding the input latent code z ∈ Z through the input layer only (Figure1-a), the style-based generator maps the input z to an intermediate latent space w ∈ W , which then controls the “style” of the image at each layer via adaptive instance normalization (AdaIN [21]) (Figure 1-b). It has been demonstrated that such a design allows a less entangled representation learning inW , leading to better generative image modeling [12, 23, 27, 40]. Despite the promising results, the style-based generator can suffer from the style control via AdaIN operation [28, 52]. Specifically, the style control is content-independent. It “washes away” the original information of features by normalization and assigns new styles decided by the latent codes regardless of the image/feature content. Besides, the style code w affects the entire image by scaling and biasing complete feature maps with a single value via AdaIN operation [35, 44, 54]. Such an imposed single style over multiple image regions can inevitably result in entangled representation of
∗This work was done while Yanhong Zeng was a research intern at Microsoft Research Asia.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
different image variations (e.g., hairstyle, facial expression) [28, 44, 52]. These limitations for image modeling can even lead to visible artifacts in the synthesized results (e.g., droplet artifacts [27]).
To get rid of the issues caused by StyleGAN’s style modeling, we introduce a new perspective that views image synthesis as a visual token generation problem. The visual token is a popular representation of an image patch with a predefined size and position [6, 7, 13]; and has shown an impressive superiority in various tasks with the development of Transformer models [6, 13, 43, 46]. Inspired by the appealing property of the token-based representation, we propose to achieve image synthesis by visual token generation. Specifically, it takes as input a sequence of latent tokens to predict the visual tokens of an image. Such a token-based representation enables a flexible local manipulation for different image regions, which makes it possible to learn content-aware and fine-grained style control for image synthesis.
Under this new paradigm, we design a token-based generator, i.e., TokenGAN, for the visual token generation problem. Specifically, we consider two different types of input tokens in the generator, i.e., the content tokens and the style tokens. The content tokens are learned as the constant input in the generator network and the style tokens are projected from a learned intermediate latent space (Figure1-c). Given a sequence of style tokens, the TokenGAN learns to control the visual token generation by rendering each content token with related style tokens according to their semantics. In particular, since the Transformer has been verified to be effective in sequence modeling in a broad range of tasks [6, 43], we adopt a generator network architecture from a visual Transformer to model the relations between the content tokens and the style tokens. Through such a content-dependent style modeling by the attention mechanism in Transformer, the TokenGAN is able to achieve content-aware and fine-grained style learning for image synthesis.
We conduct both quantitative comparisons and qualitative analysis on several unconditional image generation benchmarks. The results show that the token-based generator has achieved comparable results to the state-of-the-art in image synthesis. We summarize our contributions as below:
• We propose a new perspective of achieving image synthesis by visual token generation. Such a token-based representation enables flexible local manipulation for different image regions, leading to a better image modeling. • We propose a token-based generator (i.e., TokenGAN) for the visual token generation. Specifically, the TokenGAN introduces the style tokens and the content tokens. It adopts a Transformer-based network for content-aware style modeling. • We show extensive experiments (quantitative and qualitative comparisons, study on style editing, image inversion and image interpolation) to verify the effectiveness of the tokenbased generator. Specifically, the token-based generator is able to synthesize high-fidelity 1024× 1024 images without any convolutions in the generator.

heading: 2 Related works

heading: 2.1 Style-based generator

text:
The distinguishing feature of the style-based generator is its unconventional generator architecture [26, 27, 28, 52]. Typically, the style-based generator consists of a mapping network f and a synthesis
network g. The mapping network is used to transform the input latent code z to an intermediate latent code w ∈ W for learning a less entangled latent space. The synthesis network follows a progressive growing design that is able to first output low-resolution images that are not affected significantly by high-resolution layers [25, 27]. Such an architecture design can lead to an automatic, unsupervised separation of high-level attributes in different layers. For example, when trained on human faces, the low-resolution layers control coarse styles (e.g., pose, identity) of images and high-resolution layers control fine styles (e.g., micro-structure, color scheme) [25].
To control the attributes of the image, the style-based generator produces styles from the intermediate latent code w by affine transforms, which then control each layer of the synthesis network via adaptive instance normalization (AdaIN) [21, 35]. The AdaIN operation removes styles from previous layers by normalizing each feature map to zero mean and unit deviation, and it assigns the new styles by scaling and biasing the complete normalized feature maps. However, it has been demonstrated such an AdaIN operation can destroy the magnitude information of features and thus results in the well-known droplet artifacts [28]. Besides, the scale-specific style disentanglement in StyleGAN can be limited due to the complex semantic components in each feature map [28, 52]. In this paper, we explore a new architecture by flattening the image as a sequence of tokens, which enables fine-grained control of the image by assigning token-wise semantic-aware styles based on the attention mechanism.

heading: 2.2 Transformer in vision

text:
The Transformer typically takes as input a sequence of vectors, called tokens [43]. The Transformer is built for sequence modeling solely on attention mechanisms over the tokens, dispensing with recurrence and convolutions entirely [5, 11, 43, 30]. Due to its great success in the field of natural language processing, an increasing number of works attempt to extend Transformer for computer vision tasks [6, 7, 8, 13, 18, 36, 55]. For example, iGPT shows promising results in image classification by pre-training a sequence Transformer with the task of auto-regressive next pixel prediction and masked pixel prediction [8]. Dosoyvitskiy et al. propose to split an image into patches and feed these patches into a standard Transformer (ViT) for image classification [13]. They show that a ViT with large-scale training can trump CNNs equipped with inductive bias. DETR is a seminal work that views object detection as a direct set prediction problem, eliminating the need for hand-crafted components and achieving impressive performance by Transformer [6]. To mitigate the issue of high computation complexity associated with long sequences caused by high-resolution images, a branch of works explores lightweight Transformer (e.g.., Deformable DETR) for vision tasks [55].
Transformer is also attracting increasing attention in low-level vision tasks [14, 24, 34, 36, 42]. Parmar et al. cast image generation as an autoregressive sequence generation problem and propose Image Transformer with a local self-attention mechanism [36]. However, Image Transformer can suffer from its quadratic computation cost and its long inference time due to auto-regressive prediction. Most existing works adopt a hybrid CNN-Transformer architecture, which consists of a CNN head for feature extraction, a Transformer encoder-decoder backbone, and a CNN tail for feature decoding [7, 46, 48]. For example, IPT fully utilizes the Transformer architecture by a large-scale pre-training and achieves promising results in several image restoration tasks [7]. In parallel work, Hudson et al. introduce bipartite structure to maintain computation of linear efficiency for long-range interactions across the image based on a convolution backbone [22]. Our token-based generator inherits the network architecture from Transformer without any convolution layers and is able to yield surprising promising results for high-resolution image synthesis.

heading: 3 Approach

text: In this section, we introduce the details of the proposed token-based generator (TokenGAN). As depicted in Figure 2, the TokenGAN takes as input two kinds of tokens to generate the visual tokens of images by a Transformer. We introduce the input tokens in Section 3.1 and the token-based generator architecture in Section 3.2, following the overall optimization in Section 3.3.

heading: 3.1 Visual tokens

text:
It’s intuitive that the visual tokens of an image contain the information of both content (i.e., semantics) and related styles. In our token-based generator, we choose to separate these two kinds of information
by a sequence of content tokens and a sequence of style tokens, so that the interactions between them can be modeled and controlled for the image synthesis. We introduce the details of each as below.
Content tokens. As depicted in Figure 2-a, the TokenGAN starts from the sequence of learned constant content tokens c ∈ Rm×(p2×d) through the input layer. Specifically, d is the dimension in terms of channels, p × p is the patch size of the content token, and m is the length of the content token sequence. To maintain the position information of each token, we add position encodings p ∈ Rm×(p2×d) to each content token following a commonly used paradigm [7, 13]. For simplicity, we rewrite the notation for the content tokens as {c1, · · · , cm}, where ci ∈ Rd.
Style tokens. Given a latent code z from the input latent space Z , the mapping network adopts several MLPs to map the input z to a set of different style tokens {s1, · · · , sn} ∈ S, where si ∈ Rd. As shown in Figure 2-b, the style tokens are paired with a set of learnable semantic embedding as a key-value structure in each style modulation layer of the TokenGAN. With the Transformer modeling, all the content tokens will match with the semantic embedding and then fetch the new style from the style tokens based on the matching results. The fetched new styles are used to control the values of the content tokens, which are finally decoded to images. Such a token-wise style control enables content-aware and fine-grained style learning for image synthesis.

heading: 3.2 Token-based generator

text:
As shown in Figure 2, the TokenGAN consists of multiple layers of style blocks and each style block consists of a style normalization layer and a style modulation layer. Specifically, the style normalization layer removes the styles from previous layers by relieving the dependence on the original statistics of input token features. At the same time, the style modulation layer assigns new styles from the current layer to the content tokens. In particular, the new styles are calculated based on the pairwise interactions between the style tokens and the content tokens by a cross-attention mechanism. We introduce more details about the style normalization, content-aware style modeling, and style modulation in our generator as below.
Style normalization. We apply style normalization on the input content tokens to remove the styles from previous layers, which is denoted as:
Norm(ci) = (ci − µ(ci)) / σ(ci), i = 1, · · · ,m, (1)
where each ci ∈ Rd denotes a content token, d denotes the dimension of each content token, and m is the length of the token sequence. Specifically, the style-based generator typically removes the styles from previous layers by Instance Normalization [41], which may destroy the information conveyed by the magnitude of the feature map relative to each other [27, 28, 41]. To avoid the above issue, we adopt LayerNorm [43] to relive the style dependence while maintaining the information of the relative magnitude following previous works [28, 29, 43].
Content-aware style modeling. After style normalization, the TokenGAN calculates new styles to control the input content tokens. As shown in the style modulation in Figure 2-b, the style tokens are paired with a set of learnable semantic embedding as a key-value structure. Given the content tokens {c1, · · · , cm} from the input and the style tokens from the mapping network {s1, · · · , sn} ∈ S, the token-wise new styles are calculated by attention mechanism:
S′ = Attention(C,K,S) = softmax( CKT√
d )S, (2)
where S′ ∈ Rm×d denotes new styles for all the content tokens, the input content tokens are packed together into a matrix C ∈ Rm×d, similarly with the style tokens as matrix S ∈ Rn×d, and the matrix K ∈ Rn×d for the learnable semantic embedding in where each row vector indicates a learned semantic. We omit the linear projection in the standard attention calculation for simplicity [43].
Through such a cross-attention mechanism, the new token-wise styles are fetched for style modulation based on the matching results. Particularly, tokens with similar semantics will have similar styles (e.g., the same color for the two eyes). We visualize the attended regions for different key-value pairs (i.e., semantic-style embedding) in Figure 3, and we can find that the learned embedding K is able to attend to meaningful image regions.
Style modulation. Different from the traditional Transformer that adds the attention results back to the input features as residual features, our attention results are used as new styles by amplifying each channel of the content tokens: C′ = C S′, (3) where each row vector in C ∈ Rm×d indicates a content token, and indicates element-wise multiplication. Such a style modulation affects the operation of subsequent embedding layers (which we implemented by fully-connected layers) and thus control the style of generated images.
Implementation details. We translate the generated visual tokens to the image by concatenation and reshaping. In practice, we follow StyleGAN2 and adopt a skip-generator architecture in TokenGAN. Such a skip generator generates images in each layer at different resolutions (e.g., from 42 to 82 ) and progressively upsamples and sums the images from the previous layer by skip connections to the next layer [28, 25]. After progressive summing, the generator takes the output in the last layer as the final results.

heading: 3.3 Overall optimization

text:
Style mixing. To further encourage the styles to localize, we employ the mixing regularization technique during training [27]. To be specific, a given percentage of images are generated using two random latent codes. We run two latent codes z1, z2 through the mapping network, and have the corresponding {s11, · · · , s1n}, {s21, · · · , s2n}. We randomly choose an inject point to mix the styles, so that a part of the final style tokens from z1 and another part from z2. This regularization technique prevents the network from assuming that a set of style tokens are correlated. A similar strategy is adopted in terms of different layers following the style-based generator [27].
Optimization objectives. We denote the generated image output by the token-based generator as: G(z) = g(f(z),C), (4)
where z ∈ Z ∼ N (0, 1) indicates the input latent code, C = {c1, · · · , cm} denotes the sequence of learned constant content tokens through the input layer, and f , g, G indicate the mapping network, the generator network and the token-based generator, respectively. Our token-based generator follows the same optimization objectives used by the style-based generator [27, 28], i.e., an non-saturating logistic adversarial loss with R1 regularization. To be specific, the generator loss is:
LG = Ez∼N (0,1)[log(1 + exp(−D(G(z))))], (5) where D is the discriminator and the discriminator loss is
LD = Ex∼Pd [log(1 + exp(−D(x)))] + Ez∼N (0,1)[log(1 + exp(D(G(z))))] +R1(ψ), (6) where x ∼ Pd denotes images from the real data, and R1(ψ) is the regularization term calculated by:
R1(ψ) = γ
2 Ex∼Pd
[ ‖∇Dψ(x)‖2 ] . (7)
Specifically, the R1 regularization term [32] is computed less frequently than the main loss function following the commonly used lazy regularization strategy, thus greatly diminishing the computational cost and the overall memory usage [27, 28, 52].

heading: 4 Experiments

heading: 4.1 Experiment Setup

text:
Dataset To evaluate the token-based generator and make fair comparisons with the SOTA approach (i.e., the style-based generator [28]), we conduct experiments on the most commonly-used public research datasets, i.e., Flickr-Faces-HQ (FFHQ) [27] and Large-scale Scene Understanding (LSUN) [47]. Specifically, FFHQ consists of 70,000 high-quality images of human face, which is able to evaluate the model’s ability to synthesize high-frequency details. We conduct experiments on FFHQ with both 256× 256 and 1024× 1024 image size following previous works [15, 22]. Besides, we choose the church set of LSUN [47] to evaluate the synthesis in terms of complex scenes.
Evaluation. In this section, we conduct both quantitative and qualitative experiments on the tokenbased generator. Specifically, we report quantitative results by three commonly-used metrics, i.e., Fréchet Inception Distance (FID) [19], Precision and Recall [29, 39]. We use FID as it has been widely used in many generation tasks as an effective perceptual metric [4, 26, 27]. In addition to the image quality assessment, we use the metrics of Precision and Recall to evaluate the distribution learned by GAN. Specifically, the Precision metric intuitively measures the quality of samples from the generated images, while recall measures the proportion of real images that is covered by the learned distribution by GAN [39].
Training details We use 8 NVIDIA V100 GPUs for training each model. We build upon the public Pytorch implementation of StyleGAN22 [28]. For fair comparisons, we inherit all of the training details and parameter settings from the configuration E setup. All models are trained and report the best results under the same training iterations. For each iteration, we set the size of mini-batch as 32. We use the Adam solver [3] with the same momentum parameters β1 = 0, β2 = 0.99 to train both the generator and discriminator. We apply R1 regularization for every 16 iterations.
2https://github.com/rosinality/stylegan2-pytorch

heading: 4.2 Unconditional image synthesis

text:
Quantitative comparisons. We report quantitative comparisons with StyleGAN2 [28] due to its state-of-the-art performance. For fair comparisons, we report the results with the same training iterations on FFHQ-256 [27], FFHQ-1024 [27] and LSUN CHURCH [47] by three objective metrics, i.e., FID, Precision and Recall in Table 1. The results show that the token-based generator has achieved the state-of-the-art results in terms of both perceived quality and distribution modeling.
Qualitative results. To demonstrate the image quality of the synthesis results for the token-based generator, we randomly sample noise z ∈ Z ∼ N (0, 1) to generate images for high-resolution images (i.e., with resolution 1024× 1024) of human face and the images of complex church scenes. The uncurated results can be found in Figure 4. It shows that the token-based generator is able to synthesize both high-frequency details of the human face and plausible structures for complex scenes.
Style editing. We study and report the results of style editing by the TokenGAN in Figure 5. All the results are obtained by editing the style token of interest and then re-synthesizing images using edited style tokens. The result shows several interesting properties of the TokenGAN. First, the TokenGAN has a similar behavior as in StyleGAN2, i.e., controlling coarse/middle/fine styles at different layers (e.g., pose/hairstyle/color scheme in 82/162 − 642/1282 − 2562 layers). Besides, the TokenGAN shows localized behaviors for the style tokens. For example, in the first row, editing the style token for the pose "globally" changes the pose while remaining the styles of other regions (e.g., hairstyle, background). In the second row, editing the style token for hair length, TokenGAN "locally" turns the hair longer while remaining other styles.
Image inversion. To better apply well-trained GANs to real-world applications, the technique of GAN inversion has attracted an increasing attention [40, 53, 1, 2]. Such an inversion technique enables real image editing by searching the most accurate latent code in the learned latent space to recover the real image [53]. It has been demonstrated that the model with better inversion results tend to learn a better image modeling of the real data [28]. Specifically, we randomly sample several real images from FFHQ-256 [28] and adopt the optimization-based inversion technique used by most works [27, 40, 53]. For fair comparisons, we use the same setting following StyleGAN2 [27]. The inversion results by the style-based generator and the token-based generator are shown in Figure 6. It shows that the token-based generator is able to reconstruct real images better. For example, StyleGAN2 tends to reconstruct the lips and the headscarf in the fifth case with the same red color,
while the token-based generator is able to reconstruct accurate colors for the lips and the headscarf respectively. This is because TokenGAN renders different regions by using different style tokens via attention mechanism, while StyleGAN2 renders them by a single style vector.
For a better quantitative comparison for image inversion, we randomly sample 1,000 real images from FFHQ-256 and report the mean absolute error (MAE, range=[0,255]) and the LPIPS distance [51] of the inversion results by StyleGAN2 and TokenGAN. The results in Table 2 show that TokenGAN is able to reconstruct significantly more accurate results with much lower MAE and shorter LPIPS distance, which is in line with the visual results in Figure 6.
Image interpolation. To further explore the property of the learned latent space by the token-based generator, we perform image interpolation by a linear interpolation in the learned intermediate latent space S. In practice, we randomly sample two noises z1, z2 and run them through the mapping network to have their corresponding sequence of style tokens s1 := {s11, · · · , s1n}, s2 := {s21, · · · , s2n}. After that we perform linear interpolation for each token by s3i = α× s1i + (1− α)× s2i , α ∈ (0, 1) and use the new style tokens s3 := {s31, · · · , s3n} to synthesize a new image. As shown in Figure 7, we show the sampled styles s1, s2 in the first column and the last column, and the interpolated styles between them. The results show that the token-based generator is able to perform perceptually smooth transition between two different styles (e.g., glasses, age, hair length, coloring, etc.)

heading: 4.3 Ablation study

text:
To study the effectiveness of each component of the token-based generator, this section presents ablation studies on FFHQ-256 [27] after 0.4 million training iterations in terms of FID [19]. Specifically, we calculate the metric 10 times with different random seed and report the average results.
The number of style tokens. To study the effectiveness of the style tokens, we conduct ablation studies by using different number of style tokens. As shown in Table 3, with the growing number of the style tokens, there are more style tokens to model the distribution of styles for different semantics and the quality of images are improved. However, when the number of style tokens grows to 64, the learning will be difficult and performance would drop.
The number of content tokens. We empirically set the number of content tokens m at different layers according to the image size for the best performance and training efficiency. We conduct ablation study from layer 642 to 2562 and denote the number of content tokens in these layers as (m64,m128,m256). The results in Table 4 show that more content tokens could provide more fine-grained control over the whole images, leading to better results in image generation.
Style normalization layers. We compare different style normalization layers in Table 5. The results show that, PixelNorm [25] and LayerNorm [43] significantly outperforms the style normalizatioin by instance normalization in the token-based generator [41]. We choose LayerNorm as the style normalization layer in the token-based generator.

heading: 5 Conclusion

text:
In this paper, we propose a novel TokenGAN showing that the token-based representation of image features and styles could enable content-aware and fine-grained style learning for image synthesis. Specifically, we use Transformer to model the interaction between content tokens and style tokens, which facilitates the perceived quality of generated images and shows promising properties in terms of style editing. In the future, we will study to extend the generative transformers to extensive applications, e.g., style transfer [20], animation generation [45], image inpainting [49], etc.
Acknowledgments and Disclosure of Funding
Funding in direct support of this work: NSF of China under Grant 61672548 and U1611461, GPUs provided by Microsoft Research Asia. Additional revenues related to this work: internship at Microsoft Research Asia.

title: Improving Visual Quality of Image Synthesis by A Token-based Generator with Transformers