publications

2022

1. NeurIPS

   Chroma-VAE: Mitigating Shortcut Learning with Generative Classifiers

   Yang, W., Kirichenko, P., Goldblum, M., and Wilson, A. G.

   In Advances in Neural Information Processing Systems,  2022.

   Abs HTML PDF Code Poster Slides

   Deep neural networks are susceptible to shortcut learning, using simple features to achieve low training loss without discovering essential semantic structure. Contrary to prior belief, we show that generative models alone are not sufficient to prevent shortcut learning, despite an incentive to recover a more comprehensive representation of the data than discriminative approaches. However, we observe that shortcuts are preferentially encoded with minimal information, a fact that generative models can exploit to mitigate shortcut learning. In particular, we propose Chroma-VAE, a two-pronged approach where a VAE classifier is initially trained to isolate the shortcut in a small latent subspace, allowing a secondary classifier to be trained on the complementary, shortcut-free latent subspace. In addition to demonstrating the efficacy of Chroma-VAE on benchmark and real-world shortcut learning tasks, our work highlights the potential for manipulating the latent space of generative classifiers to isolate or interpret specific correlations.

2020

1. NeurIPS

   Incorporating Interpretable Output Constraints in Bayesian Neural Networks

   Yang, W., Lorch, L., Graule, M. A., Lakkaraju, H., and Doshi-Velez, F.

   In Advances in Neural Information Processing Systems,  2020.

   [Accepted as spotlight paper, top ~3% of papers.]

   Abs HTML PDF Code Poster Slides

   Domains where supervised models are deployed often come with task-specific constraints, such as prior expert knowledge on the ground-truth function, or desiderata like safety and fairness. We introduce a novel probabilistic framework for reasoning with such constraints and formulate a prior that enables us to effectively incorporate them into Bayesian neural networks (BNNs), including a variant that can be amortized over tasks. The resulting Output-Constrained BNN (OC-BNN) is fully consistent with the Bayesian framework for uncertainty quantification and is amenable to black-box inference. Unlike typical BNN inference in uninterpretable parameter space, OC-BNNs widen the range of functional knowledge that can be incorporated, especially for model users without expertise in machine learning. We demonstrate the efficacy of OC-BNNs on real-world datasets, spanning multiple domains such as healthcare, criminal justice, and credit scoring.

2. Thesis

   Making Decisions Under High Stakes: Trustworthy and Expressive Bayesian Deep Learning

   Yang, W.

   Senior Thesis, Harvard University,  2020.

   (PDF below contains post hoc corrections for minor errata.)

   Abs HTML PDF

   Machine learning models applied to high-stakes domains must navigate the tradeoff between (i) having enough representation power to learn effectively from high-dimensional, high-volume datasets, while (ii) avoiding cost-prohibitive errors due to model overconfidence and poor generalization. This need has led to growing interest in Bayesian neural networks (BNN), models that perform Bayesian inference on deep neural networks. BNNs are able to quantify predictive uncertainty over an inherently rich hypothesis space, which is crucial for high-stakes decision-making. In this thesis, we present two contributions that tackle the shortcomings of BNNs. First, BNN priors are defined in uninterpretable parameter space, which makes it difficult for end users to express functional prior knowledge independent of training data. We formulate two novel priors that incorporate functional constraints (i.e. what values the output should hold for any given input) that can easily be specified by end users. The resulting model is amenable to black-box inference. We demonstrate its efficacy on two high-stakes domains: (i) enforcing physiologically feasible interventions on a clinical action prediction task, and (ii) enforcing racial fairness constraints on a recidivism prediction task where the training data is biased. Next, variational approximations that are typically used for BNN posterior inference do not come with provable error guarantees, making it difficult to trust their predictive estimates. By exploiting the functional form of BNNs, we bound the predictive mean error of such approximations via maximum mean discrepancy on a reproducing kernel Hilbert space. Our bound is easily estimable and directly useful to end users as it is specified in predictive space.

3. PLOS CB

   Discovery of Hierarchical Representations for Efficient Planning

   Tomov, M. S., Yagati, S., Kumar, A., Yang, W., and Gershman, S. J.

   PLOS Computational Biology,  2020.

   Abs HTML PDF Supp

   We propose that humans spontaneously organize environments into clusters of states that support hierarchical planning, enabling them to tackle challenging problems by breaking them down into sub-problems at various levels of abstraction. People constantly rely on such hierarchical presentations to accomplish tasks big and small—from planning one’s day, to organizing a wedding, to getting a PhD—often succeeding on the very first attempt. We formalize a Bayesian model of hierarchy discovery that explains how humans discover such useful abstractions. Building on principles developed in structure learning and robotics, the model predicts that hierarchy discovery should be sensitive to the topological structure, reward distribution, and distribution of tasks in the environment. In five simulations, we show that the model accounts for previously reported effects of environment structure on planning behavior, such as detection of bottleneck states and transitions. We then test the novel predictions of the model in eight behavioral experiments, demonstrating how the distribution of tasks and rewards can influence planning behavior via the discovered hierarchy, sometimes facilitating and sometimes hindering performance. We find evidence that the hierarchy discovery process unfolds incrementally across trials. Finally, we propose how hierarchy discovery and hierarchical planning might be implemented in the brain. Together, these findings present an important advance in our understanding of how the brain might use Bayesian inference to discover and exploit the hidden hierarchical structure of the environment.

2019

1. ICML Workshop

   Output-Constrained Bayesian Neural Networks

   Yang, W.*, Lorch, L.*, Graule, M. A.*, Srinivasan, S., Suresh, A., Yao, J., Pradier, M. F., and Doshi-Velez, F.

   In 36th International Conference on Machine Learning Workshop on Uncertainty and Robustness in Deep Learning,  2019.

   Abs HTML PDF Poster Slides

   Bayesian neural network (BNN) priors are defined in parameter space, making it hard to encode prior knowledge expressed in function space. We formulate a prior that incorporates functional constraints about what the output can or cannot be in regions of the input space. Output-Constrained BNNs (OC-BNN) represent an interpretable approach of enforcing a range of constraints, fully consistent with the Bayesian framework and amenable to black-box inference. We demonstrate how OC-BNNs improve model robustness and prevent the prediction of infeasible outputs in two real-world applications of healthcare and robotics.