PAIR - AI Exploreables

Mapping LLMs with Sparse Autoencoders

By Nada Hussein, Shivam Raval, Emily Reif, Jimbo Wilson, Ari Alberich, Neel Nanda, Lucas Dixon, and Nithum Thain

October 2025

Sparse Autoencoders (SAEs) are gaining recognition for their ability to enhance the interpretability of machine learning models. By extracting understandable features, SAEs offer new ways to understand and influence model behavior. In this post, we explain what SAEs are and how to train them, how they can be used to build a global feature map for an LLM, and how you can use them to steer model behaviour.

Understanding Large Language Models

Large language models (LLMs) consist of an input layer, a series of hidden layers, and an output layer.

As a model is processing an input sentence, the output of each hidden layer is referred to as an activation.

Interpretability researchers are interested in understanding how and where a model encodes certain features in its activations.

Features represent concepts that the model might use as part of its internal reasoning to arrive at an answer. Ideally, activations would be monosemantic, with each index encoding exactly one feature. This would allow us to more easily figure out the meaning of each feature by examining the inputs for which it is activated.

For example, if the model is trying to answer the question about A Midsummer Night's Dream, it might use concepts like Shakespeare, literature, and nature.

In practice, however, each index might activate for many different features, resulting in polysemantic activations. This behaviour is reflected in the superposition hypothesis, which theorizes that models often need to represent more features than they have neurons. This makes it very difficult to disentangle the activations to understand how the model is representing features and examine its underlying reasoning mechanisms.

Training Sparse Autoencoders

The goal of an SAE is to disentangle these polysemantic activations, separating them into monesmantic vectors that we can connect with individual features. To do this, SAEs use a specific type of neural network, called an autoencoder, which is tasked to try to modify and then reconstruct activations.

Once trained, an SAE runs polysemantic activations, like the ones extracted from the model, through an encoder...

...and outputs a larger, sparser representation of the activation vector, referred to as the latent. This latent is larger to reflect the superposition hypothesis that there may be more features being represented than neurons in the activation.

To encourage monosemanticity, the autoencoder network is trained with a sparsity penalty which rewards the model for activating fewer latent neurons.

The latent can then be run through a decoder...

...to reconstruct the original activations.

To ensure that the intermediate vector is a suitably accurate decomposition of the activations, the SAE is trained to recreate the original activations by minimizing the reconstruction loss between the original activations and the reconstructions.

Labeling the Latent

Once we have an accurate and sparse decomposition, we can use an autolabeler to assign feature labels to each index of the latent. Hover over each activation index to see the top activations that fired for that label. In the Labeling Neurons section below, we will take a deeper dive into a real world example of an SAE and explore the rich labels that are extracted.

Steering with SAEs

With labels for each feature in the model, we can use the SAE to attempt to steer the model. Since we know which indices in the latent correspond to which feature, if we want the model to lean more heavily into the concept that feature represents, we can increase the activation for it in the latent. On the other hand if we want the model to reflect that concept less, we can decrease the activation for it in the latent. Once this is done, we can use the remainder of the autoencoder network to build the reconstruction and substitute it for our original activation.

For example, here we activate the cats feature and see how this changes our model behavior.

Once we've modified the latent by strengthening this feature, we can run this modified representation through the decoder…

...and create a reconstruction that reflects the steered feature. Notice that because we have modified the latent, this reconstruction is different from the original activation vector, as it now amplifies the “cat” concept.

We then plug this reconstruction back into the model, and see the model's steered response!

Use the buttons below to change the steered feature.

Labeling Neurons

Training an autoencoder lets us transform polysemantic activations into a sparse feature representation, but how do we know what each feature represents? To address this, we have to label each neuron in the latent with its corresponding feature.

A common approach is to look at the training data examples where the corresponding latent neuron activates. After obtaining this list of examples, we can pass them through a standard LLM and ask it to identify a common label to represent them.

Click an unlabeled neuron below to walk through this process on an illustrative dataset.

Activations map

After labeling the latent activations with their corresponding feature labels, we can obtain a global feature map of our model. In short, what features does the model use to reason and how do they cluster? It can be hard to sift through thousands of features, so we visualize them using UMAP.

Below, we show the map for the 16384 features of Gemma Scope. The examples are colored by hierarchical clustering, and the descriptive cluster labels are also created using an LLM. We can see that a large number of the features are devoted to programming and code, likely reflecting the training data distribution.

Hover over a feature to see its label, or search to see if a given feature is present in Gemma Scope.

Conclusion

The Gemma Scope Neuronpedia demo below lets you interact directly with the Gemma Scope SAE. Try steering the model with different features to explore how the model behaviour changes. What happens if you change the strength of steering? While there is still a lot that we don’t understand about SAEs, they are part of the important toolbox we are developing to better understand and control how LLMs work for us.

Credits

Thanks to Fernanda Viégas, Mike Mozer, James Wexler, Ryan Mullins and Martin Wattenberg for their help with this piece.

More Explorables

Can Large Language Models Explain Their Internal Mechanisms?

By Nada Hussein, Asma Ghandeharioun, Ryan Mullins, Emily Reif, Jimbo Wilson, Nithum Thain and Lucas Dixon

July 2024

The “mind” of a Large Language Model (LLM) consists of layers of interconnected artificial neurons. These layers communicate with vectors of numbers, often called the hidden representations. Deriving human-comprehensible meaning from such internals of AI systems is the focus of research into machine-learning interpretability, which has been making many exciting advances recently.

This Explorable is about one of these advances: a new family of interpretability methods called Patchscopes. The idea is to perform a kind of surgery on the neurons of an LLM, cutting out and replacing hidden representations between different prompts and layers. The key concept is the inspection prompt, which acts as a lens into the mind of an LLM, allowing the model itself to help uncover human-interpretable meaning.

Patchscopes is built on an understanding of LLMs and the transformer architecture. For a deeper dive into transformers and the schematic diagrams we use throughout this Explorable, see Appendix A.

Letting representations “talk” with Patchscopes

The Patchscopes framework leverages a simple premise: LLMs have the inherent ability to translate their own seemingly inscrutable hidden representations into human understandable text. Patching hidden representations between locations during inference allows us to inspect the information within a hidden representation, understand LLM behavior, or even augment the LLM’s behavior to improve its performance. Let’s explore this concretely with a step-by-step example.

Patchscopes is a surprisingly versatile tool. It can be used for tasks like extracting factual information, such as the country in which Barcelona is located; understanding the mechanisms of model refusal, such as refusing to respond to a prompt that is perceived to come from a person with malicious intent; debugging incorrect outputs; or even finding latent harmful information that might exist in a model even if it isn’t verbalized.

In the following sections we explore three case studies in which Patchscopes is used to: (1) discover how models resolve entities in their early layers; (2) evaluate how accurately the model’s hidden representation captures well-known concepts; and (3) augment the model’s processing of complex questions. These case studies touch upon some of the aspects of Patchscopes that require human judgment, such as determining good target locations for a given source location, or how to construct prompts that enable model inspection—although there are many more dimensions and design considerations to be explored. (For a formal definition of Patchscopes, see Appendix B.)

Patchscopes is being actively developed, and we invite your explorations, commentary, and feedback on GitHub.

Case-study: Verbalizing the model’s “thought” process

An LLM’s ability to perform a task—answering questions, summarizing documents, translating languages—is dependent on its ability to correctly contextualize the tokens in its prompt. For example, to answer the question “When was Diana, Princess of Wales born?” the model must understand that “Diana” refers to “Princess Diana”, rather than the generic name, or the ancient Roman goddess of the hunt. At what point does the model make this association, if at all?

In this case study, we use Patchscopes to shed light on how models process tokens across the layers of the model, and update their hidden representations to resolve what an entity refers to. We focus on understanding hidden representations of named entities—people, places, movies, etc.

Entity resolution & establishing context from a prompt

We are first going to identify how the model resolves the entity Diana from a source prompt. This is the first step in understanding how the model stores and accesses factual information, which can enable other tasks such as correcting outdated information. Building on the Princess Diana example, we construct a simple source prompt using her name and title:

We can measure how an entity is resolved by having the model generate a description of the entity and compare it to a known factual description. We can trigger this generation with a few-shot inspection prompt consisting of entity names followed by their description, and a final placeholder token x, which will be used as the target patching location.

We feed the source prompt into a 40 layer model and extract the hidden representation of the Wales token at every layer. Next, each hidden representation is patched into the first layer in place of the x token in the inspection prompt. Finally, the LLM continues decoding on this patched inspection prompt to generate the entity description for the given hidden representation of the Wales token.

In the table below, we show the generated descriptions from the first 10 source layers, and use an LLM-based automated evaluator to score the similarity of each output to the Wikipedia description of the original source prompt on a scale of 0-10.

Looking across all 40 layers, we start to see a pattern in the per-layer automated evaluator scores – that entity resolution typically happens in the early layers (< L=20) of the model. The general pattern of resolving in the early layers corresponds with theories about layer function: the role of early layers is to establish context from the prompt.

Tokenization can change how each layer processes information

Tokenization differs between model families, and has a strong impact on how the model navigates its embedding space. In the following visualization, we’ll focus on a tokenized representation of an example source prompt, “Dubai”:

Using the same few-shot entity description inspection prompt seen above, this example demonstrates the influence of token-by-layer processing on the generated text – we see the outputs change as the model incorporates more tokens into its context.

From TV shows to countries: Exploring classes of named entities

The visualization below allows you to explore more examples of named entities on your own. While this dataset is by no means exhaustively descriptive of the model’s behavior, we do see differences in where the model resolves the target entity across different classes—TV shows, people, countries—and even across token counts. Careful observers will notice a pattern in the way the model processes information: it builds up context by sequentially processing tokens in reverse order, gradually incorporating each token into the hidden representation. Further work is required to understand these patterns at scale.

To quickly recap this case study, we see that Patchscopes provides a highly flexible method for defining experiments that extract, verify, and characterize the model’s information retrieval process. It demonstrates highly expressive generations in early layers, where models are believed to do this processing.

Case-study: Extracting latent attributes

We can also use Patchscopes to take this a step further, exploring how a model relates subjects to their specific attributes at various model layers.

Feature extraction to corroborate hidden representations

Consider the subject Spain. Can we determine whether our LLM is able to gather context about Spain and correctly identify attributes like its largest city or official currency? More generally, given a hidden representation of a subject, we can explore whether an LLM can extract a specific attribute of that subject.

With Patchscopes, we can formulate an inspection prompt to extract an attribute. This prompt consists of a description of the attribute and a placeholder token, x, into which the hidden representation of the source subject is patched. For example, if we’re interested in the feature “the largest city” of a certain country, an inspection prompt that should generate the desired attribute in text looks like:

To extract the feature, we first capture the hidden representation from the last token of the source prompt and patch it into the x token in the inspection prompt above. We then determine if the correct answer, in this case Madrid, appears in the target model’s output. The visualization below lets you explore this example for a range of source layers.

Comparing Patchscopes to probing

The most comparable approach to this problem is a technique called probing – training a classifier (such as logistic regression) on a subset of data to allow it to predict an attribute from a hidden representation. This approach has a few downsides – namely, it requires a pre-specified subset of labels, which can limit expressivity of the output, and requires dedicated data to train the classifier.

Unlike probing, Patchscopes does not require any labeled data or supervised training. Additionally, we show that Patchscopes is more accurate than probing in many contexts in the full paper. The visualization below lets us explore a variety of reasoning tasks and compare the performance of patching and probing.

A few noticeable patterns emerge from these examples. Patchscopes outperforms probing in early layers on 7 of the 8 tasks. However, in 4 of the 8 tasks, probing outperforms patching in some later layers.

Importantly, these results suggest that Patchscopes is a viable alternative to probing classifiers, while offering two improvements. First, the flexibility and simplicity of construction for feature extraction inspection prompts indicate that the prompts could be easily generated from existing, trusted structured data that enables direct performance comparison, such as Wikipedia’s RDF dumps. Second, it does not require any additional training data or a predefined set of labels.

Application: Correcting reasoning errors

In the case studies above, we demonstrate that Patchscopes is a flexible and expressive tool for inspecting a model. Can we go further and use the same framework to change model behavior to improve its performance?

Consider multi-hop reasoning, a problem formulation where the answer depends on making logical connections between disjointed pieces of information. For example, to determine the answer to ‘the largest city in sushi’s country of origin’, a model needs to correctly recognize that Sushi’s country of origin is Japan in order to answer with Japan’s largest city, Tokyo.

One possible performance inhibitor on these problems may be the model failing to conduct sequential reasoning in the right order. Chain-of-thought (CoT) style prompts improve performance on multi-hop reasoning problems by explicitly expressing reasoning as a series of steps during generation, but there is still room for improvement. Patchscopes may provide an alternative mechanism to control the order of reasoning steps and ultimately correct the generated output.

Patching backwards

To explore this, we constructed a small dataset of two-clause multi-hop reasoning queries. We used examples from the attribute extraction case study, where the first clause defines the scope of the answer, and the second clause establishes the necessary context. The model can correctly answer each of these clauses independently, but fails to answer the composite multi-hop query. With some knowledge about the query structure, we can create a Patchscope that works backward, patching to an earlier token in the same prompt to intervene and correct its answer, as shown in the example below:

By generalizing the query structure that defines the Patchscope above, we start to see some patterns across our multi-hop reasoning queries:

Explore the accuracy of patched generations

The visualization below allows you to explore these examples individually. Use the accuracy grid in the center to explore the patched generations for different source and target layer pairs, and compare them to the model’s baseline response. As described above, the hidden representation is extracted from the last token in the multi-hop query, and we show the associated entity description below the accuracy grid to help you locate useful source layers.

The findings presented above are by no means conclusive, but they do demonstrate the utility of Patchscopes as a method for making sense of model behavior regardless of whether that model is generating the “correct” output.

Discussion and open questions

In the explorations above, we introduced Patchscopes—a method for understanding hidden representations in models—and established it as a flexible, highly expressive tool for understanding and even augmenting model behavior, which improves and unifies prior work (e.g., vocabulary projection, probing classifiers, and computational interventions) into a shared theoretical framework. We are excited that the core idea of Patchscopes has already been adopted more broadly in interpretability research as a stepping stone to advance our work on model understanding.

The case studies explored in this post capture a range of scenarios from input processing to attribute extraction to factual reasoning correction, but LLMs are used for an incredibly wide array of tasks, including mathematical reasoning, classification, code completion, and more. Each of these tasks provide many prompts to analyze, but significant research is required to understand how to create the relevant source or inspection prompts that enable human inspection of the model’s reasoning processes in these tasks. While we do not yet have generalizable heuristics or design guidelines, we do expect that certain classes of inspection prompts will have wide applicability as intermediate inspection tools across tasks that support the creation of task-specific inspection prompts. Examples of these include, but are not limited to, the few-shot entity description and zero-shot attribute extraction prompts described in this post.

Our overarching objective is to develop guidelines and tools that enable the effective application of Patchscopes across the widest possible variety of modeling tasks. To this end, we are exploring several research directions, such as broadening task applications, automating task-specific and task-agnostic patching configurations, and exploring non-identity hidden representation transforms.

The effectiveness of a patching configuration—the layer, prompt, and token position choices—depends on how the information propagates during inference. In the experiments we discussed above, simple heuristics worked well in configuring an effective Patchscope. When using the same source and target model, picking the same layer for source and target is a good start. If the analytical goal is more expressive, open-ended generation, e.g., to enable verification, we may pick target layers earlier than the source layer, as we did for entity resolution. Patching into later target layers seems to be less useful generally, as shown in the multi-hop reasoning, likely because the model has already shifted from sense-making to decoding. Considering token positions, a rule of thumb is that picking late target positions minimizes the chances of placeholder contamination, and therefore is more likely to work than other alternatives. If we aim to ask more complex questions, and therefore use more complicated inspection prompts, prior knowledge about the model’s information flow may be required, or we might need to test many configuration options before finding one that is most effective. Coming up with effective Patchscopes configurations automatically would make this framework much more powerful, and is a research direction we are currently exploring.

Finally, a full Patchscopes configuration allows for the application of transforms to the hidden representation between the extraction and injection steps. The examples in this explorable only employ the identity function as the transform, but the paper shows that, surprisingly, it is indeed possible to use a larger and more expressive model from the same family to explain a smaller model using an affine transformation function.

But what about different hidden representation dimensions, models from different families, or even models trained on different modalities of data? With more complex transformations, how far can we push the limits in cross-model Patchscopes? Answering these questions requires more research, and we continue to explore these ideas.

We encourage you to check out our paper or reach out to us on GitHub for more information about Patchscopes.

Credits

Thanks to Avi Caciularu, Carter Blum, Mor Geva, Mahima Pushkarna, Ian Tenney, Fernanda Viegas, Martin Wattenberg, David Weinberger, and James Wexler for their help with this piece.

Appendix A: A brief review of transformers

LLMs are transformer models that take a text string as input, called a prompt, that describes a task, such as answering a question, summarizing an article, translating text, generating code snippets. The LLM’s job is to generate the text that best follows from the task captured in the prompt. For an illustrative example, consider the prompt: “United Kingdom”.

This input text is tokenized – broken up into the model’s atomic vocabulary consisting of individual words and small sequences of characters.

Each token is embedded into a high-dimensional numeric vector to facilitate computation. We represent this operation with a trapezoid. The embedded vector is the first hidden representation in the model.

After embedding, LLMs are organized into layers of transformer blocks. Each layer produces an updated hidden representation based on the output of the preceding layer, which we visualize with a circle labeled $\ell^i$ , where $i$ stands for the layer index. The hidden representation produced by the first transformer layer is labeled $\ell^0$ .

This procedure continues through the final layer, gradually updating the hidden representations to incorporate context about the task described in the prompt.

Next, subsequent prompt tokens are fed into the network. Tokens are stored in an ordered sequence, and the hidden representation produced by a layer for a given token is influenced by the hidden representation for that token at the preceding layer as well as the hidden representations for all preceding tokens in the sequence. Networks that only look backwards like this are called decoder-only models and are our focus here.

After the final prompt token is processed, the final layer’s hidden representation is used to yield an output token based on the distribution of the model’s training data. This output token, is, becomes the next input to the network.

By feeding generated tokens back in as input, the model can generate phrases, sentences, and even long-form text.

For a more in-depth, visual exploration of transformers and the attention mechanism, see the related YouTube videos from 3Blue1Brown.

Appendix B: Formal description of Patchscopes

Given a hidden representation obtained from an LLM inference pass, a Patchscope instance decodes specific information from it by patching it into a different inference pass (of the same or a different LLM) that encourages the translation of that specific information.

Formally, given an input sequence of $n$ tokens $S = \langle s_1, ..., s_{n} \rangle$ and a model $M$ with $L$ layers, $\bm{h}_{i}^{\ell}$ denotes the hidden representation obtained at layer $\ell \in [1, \ldots, L]$ and position $i \in [1, \ldots, n]$ , when running $M$ on $S$ . To inspect $\bm{h}_{i}^{\ell}$ , we consider a separate inference pass of a model $M^*$ with $L^*$ layers on a target sequence $T = \langle t_1, \ldots, t_{m} \rangle$ of $m$ tokens. Specifically, we choose a hidden representation $\bar{\bm{h}}_{i^*}^{\ell^*}$ at layer $\ell^* \in [1, \ldots, L^*]$ and position ${i^*} \in [1, \ldots, m]$ in the execution of $M^*$ on $T$ . Moreover, we define a mapping function $f({\bm{h}}; \bm{\theta}): \mathbb{R}^{d} \mapsto \mathbb{R}^{d^*}$ parameterized by $\bm{\theta}$ that operates on hidden representations of $M$ , where $d$ and $d^*$ denote the hidden dimension of representations in $M$ and $M^*$ , respectively. This function can be the identity function, a linear or affine function learned on task-specific pairs of representations, or even more complex functions that incorporate other sources of data.

The patching operation refers to dynamically replacing the representation $\bar{\bm{h}}_{i^*}^{\ell^*}$ during the inference of $M^*$ on $T$ with $f(\bm{h}_{i}^{\ell})$ . Namely, by applying ${\bar{\bm{h}}}_{i^*}^{\ell^*} \leftarrow f(\bm{h}_{i}^{\ell})$ , we intervene on the generation process and modify the computation after layer $\ell^*$ .

Overall, a Patchscope intervention applied to a representation determined by $(S,i,M,\ell)$ , is defined by a quintuplet $(T,i^*,f, M^*,\ell^*)$ of an inspection prompt $T$ , a target position $i^*$ in this prompt, a mapping function $f$ , a target model $M^*$ , and a target layer $\ell^*$ of this model. It is possible that $M$ and $M^*$ are the same model, $S$ and $T$ are the same prompt, and $f$ is the identity function $\mathbb{I}$ (i.e., $\mathbb{I}(\bm{h}) = \bm{h}$ ).

Footnotes

Equal contribution.

For example, Anthropic and OpenAI recently published work exploring how to disentangle activation vectors and think about using them to control model behavior, and Google’s Responsible Generative AI Toolkit includes an interpretability-based prompt-debugger to help identify mistakes in the prompts fed to LLMs.

To score each generation, the following prompt was used with PaLM 2 Text Unicorn as the model: How semantically similar are the following texts: <text_A> {description} </text_A> <text_B> {generation} </text_B> Both these texts try to explain the following entity: <entity> {entity} </entity> Provide an integer rating between 1 and 10. 1 refers to ‘not similar at all’, and 10 refers to ‘extremely similar’: <label>

The full datasets are available on Github. While some of these examples are contrived to ease research evaluations, they do demonstrate the principle that rerouting hidden representations allows the model to generate more accurate conclusions.

Note that the model can sometimes output incorrect information about the entity it is describing.

To get a more robust estimate of accuracy for patching, we average across five different source prompts for each example. A patching instance was considered correct if the correct answer appeared in any of the first 20 generated outputs of any target layer.

As we saw in the prior case study, information about the input is more readily accessible in early to mid layers, before the model shifts toward next-token prediction. This explains why the performance of this simple Patchscope can drop for later source layers. In the open questions section, we will discuss some strategies to improve attribute extraction accuracy in later layers by adjusting the Patchscope configurations.

We want to caution that our primary goal in this section is not to devise a new method to solve multi-hop queries that is necessarily a competitor to CoT, but rather to make a proof-of-concept while comparing with CoT as a common reference.

A backward-acting Patchscope extracts the hidden representation from one token and injects it into a preceding token in the prompt, thus this is only possible when the source and inspection prompts are the same. While the exact implications of patch directionality are unknown, we believe that this technique will be useful for understanding the information flow inside the model and may aid in identifying reasoning circuits, for example.

To generate the entity descriptions, we use the same inspection prompt as seen in the entity description case study: Syria: Country in the Middle East, Leonardo DiCaprio: American actor, Samsung: South Korean multinational major appliance and consumer electronics corporation, x For each multi-hop query, we patch the hidden representation at the last token location of the selected source layer in place of the ‘x’ token of the inspection prompt at the last target layer. This generates the entity description of the hidden representation – here we are looking for some description of the correct country of origin as a checkpoint to determine whether the model will be able to answer the multi-hop query.

Sometimes the remaining representation from the placeholder token in the early layers interferes with future token generations. We call this phenomenon placeholder contamination. It is only relevant if one is interested in generating more than the next immediate token, and it is more likely to happen if target location is in the later layers.

Mechanistic interpretability work that focuses on circuits is one way of understanding information flow in a model. Prior work related to indirect object identification and factual associations are some concrete examples of how to provide precise localization guidelines for patching configurations. Conversely, we expect that investigations that use Patchscopes may result in the identification of new circuits in LLMs.

Tokenization strategies are complex and highly variable. Models are trained to work with specific tokenizers, and tokenizations may be incompatible across models.

The models we explore in this post use a deterministic embedding process to transform the token vector into the high-dimensional space of hidden representations. Most recent LLMs—including Gemma, Llama, and Mistral—use a single embedding layer for this task.

Technically speaking, transformer models don’t just operate on prompt tokens one at a time. Rather, they process prompt tokens in parallel for efficiency. We visualize them here as entering sequentially to reinforce the idea that in decoder-only models, the processing of later tokens can only be influenced by earlier tokens.

For simplicity, this post only covers case studies where the source and target model are the same. However, they can be different, and we briefly discuss in the open questions section.

Patchscopes supports applying arbitrary transformations to the hidden representation from the source location before it is injected into the target location. The examples in the post use the identity function as this transform (i.e., does nothing), but complex transforms may be necessary when patching across different models, or inspecting more complex functions.

References

Understanding intermediate layers using linear classifier probes Alain, G. and Bengio, Y. , 5th International Conference on Learning Representations, Workshop Track Proceedings, 2017.

Probing Classifiers: Promises, Shortcomings, and Advances Belinkov, Y. Computational Linguistics, 48(1):207–219, 2022.

Eliciting Latent Predictions from Transformers with the Tuned Lens Belrose, N., Furman, Z., Smith, L., Halawi, D., Ostrovsky, I., McKinney, L., Biderman, S., and Steinhardt, J., arXiv preprint arXiv:2303.08112, 2023.

Language Models are Few-Shot Learners Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., Amodei, D., arXiv preprint arXiv:2005:14165, 2020.

SelfIE: Self-Interpretation of Large Language Model Embeddings. Chen, H., Vondrick, C., & Mao, C.. arXiv preprint arXiv:2403.10949. 2024.

Patchscopes: A Unifying Framework for Inspecting Hidden Representations of Language Models Ghandeharioun, A. $^*$ , Caciularu, A. $^*$ , Pearce, A., Dixon, L., Geva, M. ICML (to appear), 2024.

Who’s asking? User personas and the mechanics of latent misalignment. Ghandeharioun, A., Yuan, A., Guerard, M., Reif, E., Lepori, M.A., and Dixon, L., arXiv preprint arXiv:2406.12094, 2024.

What do tokens know about their characters and how do they know it? Kaushal, A., Mahowald, K. 2022. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2487–2507, Seattle, United States. Association for Computational Linguistics.

Locating and editing factual associations in GPT Meng, K., Bau, D., Andonian, A., and Belinkov, Y. Advances in Neural Information Processing Systems, 35:17359–17372, 2022.

Introducing Meta Llama 3: The most capable openly available LLM to date. Meta, 2023.

Interpreting GPT: the logit lens. nostalgebraist, LessWrong, 2020.

XSTest: A Test Suite for Identifying Exaggerated Safety Behaviours in Large Language Models. Röttger, P., Kirk, H.R., Vidgen, B., Attanasio, G., Bianchi, F., and Hovy, D., arXiv preprint arXiv:2308.01263, 2023.

Attention in transformers, visually explained Sanderson, G., 3Blue1Brown, 2024.

But what is a GPT? Visual intro to transformers Sanderson, G., 3Blue1Brown, 2024.

BERT Rediscovers the Classical NLP Pipeline Tenney, I., Das, D., Pavlick, E. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4593-4601. 2019.

Attention is all you need Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., … & Polosukhin, I. Advances in neural information processing systems, 30. 2017.

The Bottom-up Evolution of Representations in the Transformer: A Study with Machine Translation and Language Modeling Objectives Voita, E., Sennrich, R., Titov, I. Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), 2019.

Chain-of-Thought Prompting Elicits Reasoning in Large Language Models. Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q.V. and Zhou, D.. Advances in Neural Information Processing Systems 35, 2022.

Judging llm-as-a-judge with mt-bench and chatbot arena. Zheng, L., Chiang, W.L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E. and Zhang, H. Advances in Neural Information Processing Systems 36, 2024.

More Explorables

Do Machine Learning Models Memorize or Generalize?

By Adam Pearce, Asma Ghandeharioun, Nada Hussein, Nithum Thain, Martin Wattenberg and Lucas Dixon

August 2023

In 2021, researchers made a striking discovery while training a series of tiny models on toy tasks . They found a set of models that suddenly flipped from memorizing their training data to correctly generalizing on unseen inputs after training for much longer. This phenomenon – where generalization seems to happen abruptly and long after fitting the training data – is called grokking and has sparked a flurry of interest .

Do more complex models also suddenly generalize after they’re trained longer? Large language models can certainly seem like they have a rich understanding of the world, but they might just be regurgitating memorized bits of the enormous amount of text they’ve been trained on . How can we tell if they’re generalizing or memorizing?

In this article we’ll examine the training dynamics of a tiny model and reverse engineer the solution it finds – and in the process provide an illustration of the exciting emerging field of mechanistic interpretability . While it isn’t yet clear how to apply these techniques to today’s largest models, starting small makes it easier to develop intuitions as we progress towards answering these critical questions about large language models.

Grokking Modular Addition

Modular addition is essentially the fruit fly of grokking. The above line chart comes from a model trained to predict $a + b \bmod 67$ . We start by randomly dividing all the $a, b$ pairs into test and training datasets. Over thousands of training steps, the training data is used to adjust the model into outputting correct answers, while the test data is only used to check if the model has learned a general solution.

The model’s architecture is similarly simple: $\text{ReLU}\left(\mathbf{a}_{\text{one-hot}} \mathbf{W}_{\text{input}} + \mathbf{b}_{\text{one-hot}} \mathbf{W}_{\text{input}}\right) \mathbf{W}_{\text{output}}$ — a one-layer MLP with 24 neurons. All the weights of the model are shown in the heatmap below; you can see how they change during training by mousing over the line chart above.

The model makes a prediction by selecting the two columns of $\mathbf{W}_{\text{input}}$ corresponding to inputs $a$ and $b$ then adding them together to create a vector of 24 separate numbers. Next it sets all the negative numbers in the vector to 0 and finally outputs the column of $\mathbf{W}_{\text{output}}$ that’s closest to the updated vector.

The weights of the model are initially quite noisy but start to exhibit periodic patterns as accuracy on the test data increases and the model switches to generalizing. By the end of training, each neuron — each row of the heatmap — cycles through high and low values several times as the input number increases from 0 to 66.

This is easier to see if we group the neurons by how often they cycle at the end of training and chart each of them as a separate line:

The periodic patterns suggest the model is learning some sort of mathematical structure; the fact that it happens when the model starts to solve the test examples hints that it’s related to the model generalizing. But why does the model move away from the memorizing solution? And what is the generalizing solution?

Generalizing With 1s and 0s

Figuring out both of these questions simultaneously is hard. Let’s make an even simpler task, one where we know what the generalizing solution should look like and try to understand why the model eventually learns it.

We’ll take random sequences of thirty 1s and 0s and train our model to predict if there is an odd number of 1s in the first three digits. e.g. 000110010110001010111001001011 is 0 while 010110010110001010111001001011 is 1 — basically a slightly trickier XOR with some distraction noise. A generalizing model should only use the first three digits of the sequence; if the model is memorizing the training data, it will also use the subsequent distracting digits .

Our model is again a one-layer MLP, trained on a fixed batch of 1,200 sequences. At first only training accuracy increases — the model is memorizing the training data. As with modular arithmetic, test accuracy is essentially random and then sharply rises as the model learns a general solution.

While memorizing , the model looks dense and noisy with lots of high magnitude weights (shown as dark red and blue squares) spread across the chart below – the model is using all the inputs to make a prediction. As the model generalizes and gets perfect test accuracy, we see all the weights connected to the distracting digits gray out with very low values and the model focusing on the first three digits — mirroring the generalized structure we expected!

With this simplified example it’s easier to see why this happens: we’re pushing our model to do two things during training — output a high probability for the correct label (called minimizing loss ) and have weights with low magnitudes (known as weight decay ). Train loss actually slightly increases before the model generalizes as it exchanges loss related to outputting the correct label for having lower weights.

The sharp drop in test loss makes it appear like the model makes a sudden shift to generalization. But if we look at the weights of the model over training, most of them smoothly interpolate between the two solutions. The rapid generalization occurs when the last weights connected to the distracting digits are pruned by weight decay.

When Does Grokking Happen?

It’s important to note that grokking is a contingent phenomenon — it goes away if model size, weight decay, data size and other hyper parameters aren’t just right. With too little weight decay, the model can’t escape overfitting the training data. Adding more weight decay pushes the model to generalize after memorizing. Increasing weight decay even more causes test and train loss to fall together; the model goes straight to generalizing. And with too much weight decay the model will fail to learn anything.

Below, we’ve trained over a thousand models on the 1s and 0s task with different hyperparameters. Training is noisy so nine models have been trained for each set of hyperparameters.

We can induce memorization and generalization on this somewhat contrived 1s and 0s task — but why does it happen with modular addition? Let’s first understand a little more about how a one-layer MLP can solve modular addition by constructing a generalizing solution that’s interpretable.

Modular Addition With Five Neurons

Recall that our modular arithmetic problem $a + b \bmod 67$ is naturally periodic, with answers wrapping around if the sum ever passes 67. Mathematically, this can be mirrored by thinking of the sum as wrapping $a$ and $b$ around a circle. The weights of the generalizing model also had periodic patterns, indicating that the solution might use this property.

We can train a simpler model with a head start on the problem, constructing an embedding matrix that places $a$ and $b$ on a circle by computing $\cos$ and $\sin$ for each possible input number $i \in \{0, 1, \ldots, 66\}$ .

Then we train $\mathbf{W}_{\text{in-proj}}$ and $\mathbf{W}_{\text{out-proj}}$ in this one-layer MLP:

With just five neurons the model finds a solution with perfect accuracy.

Eyeballing the trained parameters, all the neurons converge to roughly equal norms. If we directly plot their $\cos$ and $\sin$ components, they’re essentially evenly distributed around a circle:

Connect the adjacent neurons on the $\mathbf{W}_{\text{in-proj}}$ circle and an intriguing pattern emerges: $\mathbf{W}_{\text{out-proj}}$ is rotating around the circle twice as fast as $\mathbf{W}_{\text{in-proj}}$ .

The details of how this solution works aren’t essential — check out Appendix A to see how the doubled rotation allows the model to map inputs like $1 + 0 \bmod 67$ and $2 + 66 \bmod 67$ to the same place — but we have found a 20 parameter construction that solves modular addition. Can we find the same algorithm hidden in the 3,216 parameter model we started with? And why does the larger model switch to the generalizing solution after memorizing?

It’s Full of Stars

Here’s the $a + b \bmod 67$ model that we started with — it’s trained from scratch with no built-in periodicity.

Unlike the constructed solution, where $\mathbf{W}_{\text{embed}}$ rotates around the circle once, this model has many different frequencies.

Below, we’ve isolated the frequencies using the discrete Fourier transform (DFT). This factors out the learned periodic patterns across inputs, leaving us with the equivalent of $\mathbf{W}_{\text{in-proj}}$ and $\mathbf{W}_{\text{out-proj}}$ from the constructed solution. For each neuron, this gives a $\cos$ and $\sin$ value for every possible periodic frequency from 1 to 33. The wave charts we show above use this to group neurons into frequencies by finding their largest $\cos$ and $\sin$ value across all frequencies.

Just like in the 1s and 0s task, weight decay encourages this representation to become much sparser as the model generalizes .

Grouping neurons by their final trained frequencies, and plotting the $\cos$ and $\sin$ components of the DFT for each neuron, we see the same star shapes from the constructed solution appear.

This trained model is using the same algorithm as our constructed solution! Below, the contribution to the output generated by the neurons in each frequency are shown and we can see them calculating $\cos\frac{2\pi (a + b) freq}{67}$ .

Notice what happens to the group of neurons with a frequency of 7 when test loss improves after the short plateau at 45,000 steps — they start to snap into a star shape and their outputs more closely approximate a wave.

To lower loss without using higher weights (which would be punished by weight decay), the model uses several frequencies, taking advantage of constructive interference. There’s nothing magical about the frequencies 4, 5, 7 and 26 — click through other training runs below to see variations of this algorithm get learned.

Open Questions

While we now have a solid understanding of the mechanisms a one-layer MLP uses to solve modular addition and why they emerge during training, there are still many interesting open questions about memorization and generalization.

Which Model Constraints Work Best?

Directly training the model visualized above — $\text{ReLU} \left(a_{\text{one-hot}}\textbf{W}_{\text{input}} + b_{\text{one-hot}}\textbf{W}_{\text{input}} \right) \textbf{W}_{\text{output}}$ — does not actually result in generalization on modular arithmetic, even with the addition of weight decay. At least one of the matrices has to be factored:

$\textbf{W}_{\text{input}} = \mathbf{W}_{\text{embed}} \mathbf{W}_{\text{in-proj}}$

$\textbf{W}_{\text{output}} = \textbf{W}_{\text{out-proj}} \textbf{W}_{\text{embed}}^{\top}$

We observed that the generalizing solution is sparse after taking the discrete Fourier transformation, but the collapsed matrices have high norms. This suggests that direct weight decay on $\textbf{W}_\text{output}$ and $\textbf{W}_{\text{input}}$ doesn’t provide the right inductive bias for the task.

Broadly speaking, weight decay does steer a wide variety of models away from memorizing their training data . Other techniques that help avoid overfitting include dropout, smaller models and even numerically unstable optimization algorithms . These approaches interact in complex, non-linear ways, making it difficult to predict a priori which will ultimately induce generalization. Collapsing $\mathbf{W}_{\text{embed}} \mathbf{W}_{\text{in-proj}}$ instead of $\textbf{W}_{\text{out-proj}} \textbf{W}_{\text{embed}}^{\top}$ , for example, helps in some setups and hurts in others:

Why Is Memorization Easier Than Generalization?

One theory: there can be many more ways to memorize a training set than there are generalizing solutions. So statistically, memorization should be more likely to happen first, especially if we have no or little regularization. Regularization techniques, like weight decay, prioritize certain solutions over others, for example, preferring “sparse” solutions over “dense” ones.

Recent work suggests that generalization is associated with well-structured representations . However, it’s not a necessary condition; some MLP variations without symmetric inputs learn less “circular” representations when solving modular addition . We also observed that well-structured representations are not a sufficient condition for generalization. This small model (trained with no weight decay) starts generalizing, then switches to memorizing with periodic embeddings.

It’s even possible to find hyperparameters where models start generalizing, then switch to memorizing, then switch back to generalizing!

What About Larger Models?

Does grokking happen in larger models trained on real world tasks? Earlier observations reported the grokking phenomenon in algorithmic tasks in small transformers and MLPs . Grokking has subsequently been found in more complex tasks involving images, text, and tabular data within certain ranges of hyperparameters . It’s also possible that the largest models, which are able to do many types of tasks, may be grokking many things at different speeds during training .

There have also been promising results in predicting grokking before it happens. Though some require knowledge of the generalizing solution or the overall data domain , some rely solely on the analysis of the training loss and might also apply to larger models — hopefully we’ll be able to build tools and techniques that can tell us when a model is parroting memorized information and when it’s using richer models.

Understanding the solution to modular addition wasn’t trivial. Do we have any hope of understanding larger models? One route forward — like our digression into the 20 parameter model and the even simpler boolean parity problem — is to: 1) train simpler models with more inductive biases and fewer moving parts, 2) use them to explain inscrutable parts of how a larger model works, 3) repeat as needed. We believe this could be a fruitful approach to better understanding larger models, and complementary to efforts that aim to use larger models to explain smaller ones and other work to disentangle internal representations . Moreover, this kind of mechanistic approach to interpretability, in time, may help identify patterns that themselves ease or automate the uncovering of algorithms learned by neural networks.

Credits

Thanks to Ardavan Saeedi, Crystal Qian, Emily Reif, Fernanda Viégas, Kathy Meier-Hellstern, Mahima Pushkarna, Minsuk Chang, Neel Nanda and Ryan Mullins for their help with this piece.

Model training code // Visualization code

Appendix A: How the Circular Construction Works

We can almost calculate $a + b \bmod 67$ using two circular embeddings and a completely linear model.

It works! But we’re cheating a bit, do you see how unembed loops around the circle twice? We need to output a single prediction for ““ — not separate predictions for ““ and ““. Directly adding the two predictions for a number together won’t work since they’re on opposite sides of the circles and will cancel each other out.

Instead, let’s incorporate a $\text{ReLU}(x)$ to fix the repeated outputs.

We’ve essentially wrapped the circle around in on itself and the model outputs a single prediction for ““.

Formally, this is the constructed model:

With modulus $M$ and $N$ evenly spaced neurons/directions:

Interestingly this circle has a few wrinkles: this construction doesn’t give an exact answer!

Neurons Modulus

Using $x^2$ instead of $\text{ReLU}(x)$ as the activation function, as suggested by gives a provably exact solution!

For simplicity, let $\omega:=\frac{2\pi}{M}$ (the angle between numbers in $\mathbf{W}_{\text{embed}}$ ) and $\theta := \frac{2\pi}{N}$ (the angle between neurons in $\mathbf{W}_{\text{in-proj}}^T$ ).

Let’s rewrite $\text{logits}^{a, b}$ as an $M$ -dimensional vector $\lVert \it{l} \rVert ^M$ where:

$l_{j} = \sum_{i=0}^{N-1} \biggl(\bigl[ \cos(a\omega-i\theta) + \cos(b\omega-i\theta) \bigl]^2\cos(j\omega-2i\theta)) \biggr)$

This follows from the $\text{logits}^{a,b}$ equation above by plugging in the definitions of $\textbf{W}_\text{in-proj}$ and $\textbf{W}_\text{out-proj}$ and applying the trigonometric identity that $\cos(x)\cos(y) + \sin(x)\sin(y) = \cos(x-y)$ .

We can then prove the following:

$\argmax_c \! \text{logits}^{a,b} = a + b \bmod M$

Applying the two trigonometric identities of $\cos(x) + \cos(y) = 2 \cos(\frac{x-y}{2}) \cos(\frac{x+y}{2})$ and $\cos^2(x)cos(y) = 1/4 \bigl[ 2\cos(y) + \cos(2x-y) + \cos (2x+y) \bigl]$ , we have:

Note that $\sum \cos(\gamma_{i})=0$ where $\gamma_{i}$ is equally spread around the circle. The first and the third sum terms wrap around the circle with $2\theta$ and $4\theta$ increments respectively. The sum of the first terms equals zero for $N \gt 2$ and the sum of the third terms equals zero for $N \gt 4$ . Therefore, we have:

$\text{logits}^{a,b} = \cos^2(\frac{a-b}{2}\omega) \cos((a+b-c)\omega)$

Since the first term is a positive constant w.r.t inputs, the equation is maximized when $\cos((a+b-c)\omega)$ is maximized, which is when $c = a + b \bmod M$ .

Essentially $\text{ReLU}(x)$ activations with weight decay (a very typical model setup) gives the model an inductive bias that’s close enough to the exact generalizing solution of $x^2$ activations with a sparse discrete Fourier transform to push in the direction of generalization but not so close that it won’t also learn to fit the training data with memorization.

Footnotes

In modular addition, we have two input numbers, $a$ and $b$ , and a modulus $m$ . We want to find the remainder of $a + b$ when divided by $m$ . This type of addition is often called clock-face addition, because when adding two times, we often report the result modulo 12 (i.e. 5 hours after 8 o’clock is 1 o’clock). Modular addition sounds simple and it is. We can easily train 1,000s of models and treat them like fruit flies in neuroscience: small enough such that it is feasible to extract their connectome synapse-by-synapse, yet providing new interesting insights about the system more broadly. We can get a good understanding of the small models we’ve trained by visualizing all their internals.

67 isn’t a magic number – we could pick many numbers to illustrate grokking, but 67 is not so small that the task is trivial and also not so large that the visualizations are overwhelming.

The model is trained with cross-entropy loss, AdamW and full batches. The section on regularization and training colab have additional details. If you’re not familiar with MLPs, playground.tensorflow.org is a great place to start. A quick notation explanation: The columns of $\mathbf{W}_{\text{input}}$ and $\mathbf{W}_{\text{ouput}}$ represent the numbers from 0 to 66. $\mathbf{a}_{\text{one-hot}}$ and $\mathbf{b}_{\text{one-hot}}$ are how we encode the model’s inputs; each pick a single column from $\mathbf{W}_{\text{input}}$ . $\text{ReLU}$ replaces negative numbers with 0s; it is a fancy) way of writing $\max(x, 0)$ .

With a small twist — we’re only outputting 1 or 0, so $\mathbf{W}_{\text{output}}$ can be a single column. In the modular addition task we needed a column for every output number. The last column of $\mathbf{W}_{\text{input}}$ is also fixed to 1 to provide a bias term.

Appendix D of “A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnetworks” has an explanation of the 4 neuron solution generalizing solution here

So far we’ve been charting accuracy, the percentage of sequences where the correct label is the most likely. Training typically instead optimizes a differentiable objective function. All the models in this post use cross entropy loss which heavily penalizes incorrect predictions with high probabilities. Note that while some formulations of loss include a weight decay or regularization term, the loss plots here depict the cross entropy component alone.

On the 1s and 0s task here, we use L1 weight decay $L1(\mathbf{w}) = \sum_{i} |w_i|$ . L2 weight decay $L2(\mathbf{w}) = \sum_{i} w_i^2$ is a more typical choice. It pushes for lots of small weights leading to redundant neurons on this task:

A model overfits the training data when it performs well on the training data but poorly on the test data — this is what we see with our memorizing models. In general, simpler models are less prone to overfitting as, due to their simplicity, decision rules are coarser and are required to make more generalizations. Of course, if a model is too simple for a task, it may not be able to learn good decision rules that capture the nuances of the task. Researchers force models to be simpler through a variety of techniques, including having models with fewer parameters or encouraging the parameters that the model does have to be small in size with weight decay.

Computing $\cos(i\frac{2\pi}{67})$ and $\sin (i\frac{2\pi}{67})$ gives us points evenly spaced around the unit circle. Here’s what $\textbf{W}_{\text{embed}}$ looks like on the unit circle:

The Discrete Fourier Transform helps analyze the periodic nature of a sequence of values (in this case the weights for a particular neuron) by breaking it down into sine and cosine functions. The more periodic a function is, the easier it is to represent with sine and cosines, and the sparser the output of the DFT.

We’ve reindexed the neurons by their final frequency and phase to make this grouping easier to see .

The model generates probabilities by taking the dot product of the neuron activations for a given input with $\textbf{W}_{\text{output}}$ and softmaxing. If we calculate the dot product using only the activations from neurons of a single frequency, we can see which outputs the frequency group is making more or less likely. Appendix A explains why these logits form a wave — each group of frequencies is essentially outputting how close the correct answer is to every number on a version of $\textbf{W}_{\text{embed}}$ with the group’s frequency.

Both of these models are quite small. The bottom model has tweaked hyperparameters to encourage eventual generalization: it’s slightly larger to allow it to exit local minimums, it has more training data (making low loss memorizing solutions harder to find) and it has weight decay.

References

Grokking: Generalization Beyond Overfitting On Small Algorithmic Datasets Power, A., Burda, Y., Edwards, H., Babuschkin, I., & Misra, V. (2022). arXiv preprint arXiv:2201.02177.

Omnigrok: Grokking Beyond Algorithmic Data Liu, Z., Michaud, E. J., & Tegmark, M. (2022, September). In The Eleventh International Conference on Learning Representations.

A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations Chughtai, B., Chan, L., Nanda, N. (2023). International Conference on Machine Learning.

The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networks Zhong, Z., Liu, Z., Tegmark, M., & Andreas, J. (2023). arXiv preprint arXiv:2306.17844.

Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, Cyril Zhang. (2022) Advances in Neural Information Processing Systems, 35, 21750-21764.

Grokking modular arithmetic Andrey Gromov (2023). arXiv preprint arXiv:2301.02679.

On the Dangers of Stochastic Parrots: Can Language Models Be Too Big?🦜 Bender, E. M., Gebru, T., McMillan-Major, A., & Shmitchell, S. (2021, March). In Proceedings of the 2021 ACM conference on fairness, accountability, and transparency (pp. 610-623).

Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task Li, K., Hopkins, A. K., Bau, D., Viégas, F., Pfister, H., & Wattenberg, M. (2022, September). In The Eleventh International Conference on Learning Representations.

Mechanistic Interpretability, Variables, and the Importance of Interpretable Bases Olah, C., 2022. Transformer Circuits Thread.

Progress Measures for Grokking via Mechanistic Interpretability Nanda, N., Chan, L., Lieberum, T., Smith, J., & Steinhardt, J. (2022, September). In The Eleventh International Conference on Learning Representations.

A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnetworks William Merrill, Nikolaos Tsilivis, Aman Shukla. (2023). arXiv preprint arXiv:2303.11873.

Unifying Grokking and Double Descent Davies, X., Langosco, L., & Krueger, D. (2022, November). In NeurIPS ML Safety Workshop.

Double Descent Demystified: Identifying, Interpreting & Ablating the Sources of a Deep Learning Puzzle Rylan Schaeffer, R., Khona, M., Robertson, Z., Boopathy, A., Pistunova, K., Rocks, J., Rani Fiete, I., & Koyejo, O. (2023). arXiv preprint arXiv:2303.14151.

The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon Thilak, V., Littwin, E., Zhai, S., Saremi, O., Paiss, R., & Susskind, J. (2022). arXiv preprint arXiv:2206.04817.

Towards Understanding Grokking: An Effective Theory of Representation Learning Liu, Z., Kitouni, O., Nolte, N. S., Michaud, E., Tegmark, M., & Williams, M. (2022). Advances in Neural Information Processing Systems, 35, 34651-34663.

The Goldilocks Zone: Towards Better Understanding of Neural Network Loss Landscapes Fort, S., & Scherlis, A. (2019, July). In Proceedings of the AAAI conference on artificial intelligence (Vol. 33, No. 01, pp. 3574-3581).

The Quantization Model of Neural Scaling Eric J. Michaud, Ziming Liu, Uzay Girit, Max Tegmark, O. (2023). arXiv preprint arXiv:2303.13506.

Grokking of Hierarchical Structure in Vanilla Transformers Murty, S., Sharma, P., Andreas, J., & Manning, C. D. (2023). arXiv preprint arXiv:2305.18741.

Predicting Grokking Long Before it Happens: A Look Into the Loss Landscape of Models Which Grok Notsawo Jr, P., Zhou, H., Pezeshki, M., Rish, I., & Dumas, G. (2023). arXiv preprint arXiv:2306.13253.

Language models can explain neurons in language models Bills, S., Cammarata, N., Mossing, D., Tillman, H., Gao, L., Goh, G., Sutskever, I., Leike, J., Wu, J., & Saunders, W. 2023. OpenAI Blog

Does Circuit Analysis Interpretability Scale? Evidence from Multiple Choice Capabilities in Chinchilla Tom Lieberum, Matthew Rahtz, János Kramár, Neel Nanda, Geoffrey Irving, Rohin Shah, Vladimir Mikulik (2023). arXiv preprint arXiv:2307.09458.

Toy Models of Superposition Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D., Wattenberg, M. and Olah, C., 2022. Transformer Circuits Thread.

The Connectome of an Insect Brain Winding, M., Pedigo, B. D., Barnes, C. L., Patsolic, H. G., Park, Y., Kazimiers, T., … & Zlatic, M. (2023). Science, 379(6636), eadd9330.

Multi-Scale Feature Learning Dynamics: Insights for Double Descent Pezeshki, M., Mitra, A., Bengio, Y., & Lajoie, G. (2022, June). In the International Conference on Machine Learning (pp. 17669-17690). PMLR.

Superposition, Memorization, and Double Descent Henighan, T., Carter, S., Hume, T., Elhage, N., Lasenby, R., Fort, S., Schiefer, N., and Olah, C., 2023. Transformer Circuits Thread.

More Explorables

What Have Language Models Learned?

By asking language models to fill in the blank, we can probe their understanding of the world.

Large language models are making it possible for computers to write stories, program a website and turn captions into images.

One of the first of these models, BERT, is trained by taking sentences, splitting them into individual words, randomly hiding some of them, and predicting what the hidden words are. After doing this millions of times, BERT has “read” enough Shakespeare to predict how this phrase usually ends:

This page is hooked up to a version of BERT trained on Wikipedia and books.¹ Try clicking on different words to see how they’d be filled in or typing in another sentence to see what else has BERT picked up on.

Cattle or Clothes?

Besides Hamlet’s existential dread, the text BERT was trained on also contains more patterns:

Cattle and horses aren’t top purchase predictions in every state, though! In New York, some of the most likely words are clothes, books and art:

There are more than 30,000 words, punctuation marks and word fragments in BERT’s vocabulary. Every time BERT fills in a hidden word, it assigns each of them a probability. By looking at how slightly different sentences shift those probabilities, we can get a glimpse at how purchasing patterns in different places are understood.

You can edit these sentences. Or try one of these comparisons to get started:

To the extent that a computer program can “know” something, what does BERT know about where you live?

What’s in a Name?

This technique can also probe what associations BERT has learned about different groups of people. For example, it predicts people named Elsie are older than people named Lauren:

It’s also learned that people named Jim have more typically masculine jobs than people named Jane:

These aren’t just spurious correlations — Elsies really are more likely to be older than Laurens. And occupations the model associates with feminine names are held by a higher percentage of women.

Should we be concerned about these correlations? BERT was trained to fill in blanks in Wikipedia articles and books — it does a great job at that! The problem is that the internal representations of language these models have learned are used for much more – by some measures, they’re the best way we have of getting computers to understand and manipulate text.

We wouldn’t hesitate to call a conversation partner or recruiter who blithely assumed that doctors are men sexist, but that’s exactly what BERT might do if heedlessly incorporated into a chatbot or HR software:

Adjusting for assumptions like this isn’t trivial. Why machine learning systems produce a given output still isn’t well understood – determining if a credit model built on top of BERT rejected a loan application because of gender discrimation might be quite difficult.

Deploying large language models at scale also risks amplifying and perpetuating today’s harmful stereotypes. When prompted with “Two Muslims walked into a…”, for example, GPT-3 typically finishes the sentence with descriptions of violence.

How Can We Fix This?

One conceptually straightforward approach: reduce unwanted correlations from the training data to mitigate model bias.

Last year a version of BERT called Zari was trained with an additional set of generated sentences. For every sentence with a gendered noun, like boy or aunt, another sentence that replaced the noun with its gender-partner was added to the training data: in addition to “The lady doth protest too much,” Zari was also trained on “The gentleman doth protest too much.”

Unlike BERT, Zari assigns nurses and doctors an equal probability of being a “she” or a “he” after being trained on the swapped sentences. This approach hasn’t removed all the gender correlations; because names weren’t swapped, Zari’s association between masculine names and doctors has only slightly decreased from BERT’s. And the retraining doesn’t change how the model understands nonbinary gender.

Something similar happened with other attempts to remove gender bias from models’ representations of words. It’s possible to mathematically define bias and perform “brain surgery” on a model to remove it, but language is steeped in gender. Large models can have billions of parameters in which to learn stereotypes — slightly different measures of bias have found the retrained models only shifted the stereotypes around to be undetectable by the initial measure.

As with other applications of machine learning, it’s helpful to focus instead on the actual harms that could occur. Tools like AllenNLP, LMdiff and the Language Interpretability Tool make it easier to interact with language models to find where they might be falling short. Once those shortcomings are spotted, task specific mitigation measures can be simpler to apply than modifying the entire model.

It’s also possible that as models grow more capable, they might be able to explain and perform some of this debiasing themselves. Instead of forcing the model to tell us the gender of “the doctor,” we could let it respond with uncertainty that’s shown to the user and controls to override assumptions.

Credits

Adam Pearce // July 2021

Thanks to Ben Wedin, Emily Reif, James Wexler, Fernanda Viégas, Ian Tenney, Kellie Webster, Kevin Robinson, Lucas Dixon, Ludovic Peran, Martin Wattenberg, Michael Terry, Tolga Bolukbasi, Vinodkumar Prabhakaran, Xuezhi Wang, Yannick Assogba, and Zan Armstrong for their help with this piece.

Footnotes

The BERT model used on this page is the Hugging Face version of bert-large-uncased-whole-word-masking. “BERT” also refers to a type of model architecture; hundreds of BERT models have been trained and published. The model and chart code used here are available on GitHub.

Notice that “1800”, “1900” and “2000” are some of the top predictions, though. People aren’t actually more likely to be born at the start of a century, but in BERT’s training corpus of books and Wikipedia articles round numbers are more common.

Comparing BERT and Zari in this interface requires carefully tracking tokens during a transition. The BERT Difference Plots colab has ideas for extensions to systemically look at differences between the models’ output.

This analysis shouldn’t stop once a model is deployed — as language and model usage shifts, it’s important to continue studying and mitigating potential harms.

Appendix: Differences Over Time

In addition to looking at how predictions for men and women are different for a given sentence, we can also chart how those differences have changed over time:

The convergence in more recent years suggests another potential mitigation technique: using a prefix to steer the model away from unwanted correlations while preserving its understanding of natural language.

Using “In $year” as the prefix is quite limited, though, as it doesn’t handle gender-neutral pronouns and potentially increases other correlations. However, it may be possible to find a better prefix that mitigates a specific type of bias with just a couple of dozen examples.

Closer examination of these differences in differences also shows there’s a limit to the facts we can pull out of BERT this way.

Below, the top row of charts shows how predicted differences in occupations between men and women change between 1908 and 2018. The rightmost chart shows the he/she difference in 1908 against the he/she difference in 2018.

The flat slope of the rightmost chart indicates that the he/she difference has decreased for each job by about the same amount. But in reality, shifts in occupation weren’t nearly so smooth and some occupations, like accounting, switched from being majority male to majority female.

This reality-prediction mismatch could be caused by lack of training data, model size or the coarseness of the probing method. There’s an immense amount of general knowledge inside of these models — with a little bit of focused training, they can even become expert trivia players.

More Explorables

Confidently Incorrect Models to Humble Ensembles

From Confidently Incorrect Models to Humble Ensembles

When a machine learning model is trained, it learns patterns from data. When models are deployed in the real world, however, they often encounter data that is different in some way than the data they were trained on. In this Explorable, we’ll examine how this phenomenon can cause models to misbehave and what can be done about it.

Let’s look at an example of an image classifier that is trained on MNIST, a dataset of hand-drawn digits.

The model will update its predictions as you draw on it. Try changing the eight to zero or the seven to a nine. As you experiment with drawings, do you always agree with the model’s decisions? What do you think of the model’s predictions for transition images between one number and the next?

Models Can Be Confidently Incorrect

Below, we expand the default set of images further to show what happens when the model is asked to classify something other than a number. As you click through the different images, you can see how the model’s prediction changes over time as one image transforms into another.

A typical machine learning model, when faced with data it wasn’t trained on (called out of distribution data), will try to make predictions based on the data it initially saw. In this case, no matter what we draw, the classifier will try to tell us that it is a number.

This is a consequence of the way that this (and indeed most) machine learning classifiers are designed, with a fixed set of classes in mind. The model is only able to choose which among these classes it thinks is more likely and assign a confidence score. To read more about how model confidence scores work and how they relate to probabilities, check out our previous Explorable Are Model Predictions Probabilities?

For many practical applications, however, we don’t want the model to report that a shoe is a “2” with 100% confidence. Instead, we might want the model to reduce the score that it assigns to data that is unlike anything it has been trained on. When a model incorrectly gives an out of distribution example a very high score, we call it confidently incorrect.

Combining Models Reduces Overconfidence

By averaging the output of multiple models, a technique known as ensembling, we can create a model that isn’t as confidently incorrect.

Deep Ensembles train multiple models with different random initializations. The members of the ensemble learn to make different decisions and typically disagree in their predictions on out-of-distribution inputs. While an individual network can be overconfident in its incorrect prediction, ensemble output averages the different predictions and is less confident overall on out-of-distribution inputs.

There are generally no constraints on a model’s predictive behavior away from the training data, so different models will settle into making overconfident, but different, predictions on out of distribution data. This has the effect of lowering the overall confidence of the prediction by using the disagreement between the individual models’ predictions.

Why Do Ensembles Work?

To better understand why ensembles reduce the chance of being confidently incorrect, let’s explore what happens in the case of 2-dimensional classification. Below, we plot an ensemble classifier that is deciding whether a point in 2 dimensions should be red or blue. The color of the background indicates the decision made by the ensemble of models.

With just a single model, the decision boundary is very steep and the model is quick to make a confident decision even if the point is close to the decision boundary. As the number of models grows, the disagreement between them manifests as a widening of the decision boundary, especially further away from the training data.

Not all models have simple linear decision boundaries. Below are models with piecewise linear decision boundaries (which are common in certain machine learning models). As you can see, because these have even more degrees of freedom than simple linear decision boundaries, the uncertainty of an ensemble is even more pronounced further from the data.

While ensembles increase uncertainty in a useful way, the way they do it is deeply dependent on model architecture. Moving a blue point to the lower right creates very different decision boundaries for the linear and piecewise linear models.

Beyond Ensembles

Ensembles are just one technique that can be used to improve the quality of a model’s uncertainty estimates. There are other methods that aggregate the predictions of a collection of models, but also strategies like incorporating lots more data. Modern models are big, so keeping multiple models around for an ensemble can be a non-starter for many memory or latency-sensitive applications. Therefore, a variety of factorized or efficient ensemble methods have been developed to create an ensemble within a single model.

You may have heard of Bayesian methods such as Gaussian processes, which use a carefully weighted average over all possible (infinitely many) models. These methods are expensive as you might imagine, but the research community is actively developing clever approximations and ways of incorporating them into just parts of the overall model to make things more tractable.

Another good approach to improving uncertainty estimates on out of distribution data is to pre-train bigger models on more data, in the hopes that this overall improves the quality of the representations and makes it easier to draw good decision boundaries.

Plex explores a range of techniques to improve uncertainty in modern large text and vision models. At that scale, full ensembles are just too big, but they tried combinations of efficient ensembles, incorporating more data and approximate Bayesian methods. One of these strategies is to change just the last layer of your neural network to a Gaussian Process. This kind of creates an infinite ensemble of just the last layer (making it tractable), but has an added benefit of allowing the decision boundaries of the classifier to much more directly take into account the distance between examples we are using and the training data.

Credits

Nithum Thain, Adam Pearce, Jasper Snoek and Balaji Lakshminarayanan // March 2023

Thanks to Nicole Mitchell, Lucas Dixon, and Alexander D’Amour for their help with this piece.

If you’d like to experiment with your own image classification models, check out our model training colab.

More Explorables

Collecting Sensitive Information

How randomized response can help collect sensitive information responsibly

Giant datasets are revealing new patterns in cancer, income inequality and other important areas. However, the widespread availability of fast computers that can cross reference public data is making it harder to collect private information without inadvertently violating people's privacy. Modern randomization techniques can help preserve anonymity.

Anonymous Data

Let's pretend we're analysts at a small college, looking at anonymous survey data about plagiarism.

We've gotten responses from the entire student body, reporting if they've ever plagiarized or not. To encourage them to respond honestly, names were not collected.

The data here has been randomly generated

On the survey students also report several bits of information about themselves, like their age...

...and what state they're from.

This additional information is critical to finding potential patterns in the data—why have so many first-years from New Hampshire plagiarized?

Revealed Information

But granular information comes with a cost.

One student has a unique age/home state combination. By searching another student database for a 19-year old from Vermont we can identify one of the plagiarists from supposedly anonymous survey data.

Increasing granularity exacerbates the problem. If the students reported slightly more about their ages by including what season they were born in, we'd be able to identify about a sixth of them.

This isn't just a hypothetical: A birthday / gender / zip code combination uniquely identifies 83% of the people in the United States.

With the spread of large datasets, it is increasingly difficult to release detailed information without inadvertently revealing someone's identity. A week of a person's location data could reveal a home and work address—possibly enough to find a name using public records.

Randomization

One solution is to randomize responses so each student has plausible deniability. This lets us buy privacy at the cost of some uncertainty in our estimation of plagiarism rates.

Step 1: Each student flips a coin and looks at it without showing anyone.

Step 2: Students who flip heads report plagiarism, even if they haven't plagiarized.

Students that flipped tails report the truth, secure with the knowledge that even if their response is linked back to their name, they can claim they flipped heads.

With a little bit of math, we can approximate the rate of plagiarism from these randomized responses. We'll skip the algebra, but doubling the reported non-plagiarism rate gives a good estimate of the actual non-plagiarism rate.

Flip coins

How far off can we be?

If we simulate this coin flipping lots of times, we can see the distribution of errors.

The estimates are close most of the time, but errors can be quite large.

Flip coins 200 times

Reducing the random noise (by reducing the number of students who flip heads) increases the accuracy of our estimate, but risks leaking information about students.

If the coin is heavily weighted towards tails, identified students can't credibly claim they reported plagiarizing because they flipped heads.

One surprising way out of this accuracy-privacy tradeoff: carefully collect information from even more people.

If we got students from other schools to fill out this survey, we could accurately measure plagiarism while protecting everyone's privacy. With enough students, we could even start comparing plagiarism across different age groups again—safely this time.

Conclusion

Aggregate statistics about private information are valuable, but can be risky to collect. We want researchers to be able to study things like the connection between demographics and health outcomes without revealing our entire medical history to our neighbors. The coin flipping technique in this article, called randomized response, makes it possible to safely study private information.

You might wonder if coin flipping is the only way to do this. It's not—differential privacy can add targeted bits of random noise to a dataset and guarantee privacy. More flexible than randomized response, the 2020 Census will use it to protect respondents' privacy. In addition to randomizing responses, differential privacy also limits the impact any one response can have on the released data.

Credits

Adam Pearce and Ellen Jiang // September 2020

Thanks to Carey Radebaugh, Fernanda Viégas, Emily Reif, Hal Abelson, Jess Holbrook, Kristen Olson, Mahima Pushkarna, Martin Wattenberg, Michael Terry, Miguel Guevara, Rebecca Salois, Yannick Assogba, Zan Armstrong and our other colleagues at Google for their help with this piece.

More Explorables

Why Some Models Leak Data

Machine learning models use large amounts of data, some of which can be sensitive. If they're not trained correctly, sometimes that data is inadvertently revealed.

Let’s take a look at a game of soccer.

Using the position of each player as training data, we can teach a model to predict which team would get to a loose ball first at each spot on the field, indicated by the color of the pixel.

It updates in real-time—drag the players around to see the model change.

This model reveals quite a lot about the data used to train it. Even without the actual positions of the players, it is simple to see where players might be.

Click this button to move the players

Take a guess at where the yellow team’s goalie is now, then check their actual position. How close were you?

Sensitive Salary Data

In this specific soccer example, being able to make educated guesses about the data a model was trained on doesn’t matter too much. But what if our data points represent something more sensitive?

We’ve fed the same numbers into the model, but now they represent salary data instead of soccer data. Building models like this is a common technique to detect discrimination. A union might test if a company is paying men and women fairly by building a salary model that takes into account years of experience. They can then publish the results to bring pressure for change or show improvement.

In this hypothetical salary study, even though no individual salaries have been published, it is easy to infer the salary of the newest male hire. And carefully cross referencing public start dates on LinkedIn with the model could almost perfectly reveal everyone’s salary.

Because the model here is so flexible (there are hundreds of square patches with independently calculated predictions) and we have so few data points (just 22 people), it is able to “memorize” individual data points. If we’re looking to share information about patterns in salaries, a simpler and more constrained model like a linear regression might be more appropriate.

By boiling down the 22 data points to two lines we’re able to see broad trends without being able to guess anyone’s salary.

Subtle Leaks

Removing complexity isn’t a complete solution though. Depending on how the data is distributed, even a simple line can inadvertently reveal information.

In this company, almost all the men started several years ago, so the slope of the line is especially sensitive to the salary of the new hire.

Is their salary higher or lower than average? Based on the line, we can make a pretty good guess.

Notice that changing the salary of someone with a more common tenure barely moves the line. In general, more typical data points are less susceptible to being leaked. This sets up a tricky trade off: we want models to learn about edge cases while being sure they haven’t memorized individual data points.

Real World Data

Models of real world data are often quite complex—this can improve accuracy, but makes them more susceptible to unexpectedly leaking information. Medical models have inadvertently revealed patients’ genetic markers. Language models have memorized credit card numbers. Faces can even be reconstructed from image models:

Fredrikson et al were able to extract the image on the left by repeatedly querying a facial recognition API. It isn’t an exact match with the individual’s actual face (on the right), but this attack only required access to the model’s predictions, not its internal state.

Protecting Private Data

Training models with differential privacy stops the training data from leaking by limiting how much the model can learn from any one data point. Differentially private models are still at the cutting edge of research, but they’re being packaged into machine learning frameworks, making them much easier to use. When it isn’t possible to train differentially private models, there are also tools that can measure how much data is the model memorizing. Also, standard techniques such as aggregation and limiting how much data a single source can contribute are still useful and usually improve the privacy of the model.

As we saw in the Collecting Sensitive Information Explorable, adding enough random noise with differential privacy to protect outliers like the new hire can increase the amount of data required to reach a good level of accuracy. Depending on the application, the constraints of differential privacy could even improve the model—for instance, not learning too much from one data point can help prevent overfitting.

Given the increasing utility of machine learning models for many real-world tasks, it’s clear that more and more systems, devices and apps will be powered, to some extent, by machine learning in the future. While standard privacy best practices developed for non-machine learning systems still apply to those with machine learning, the introduction of machine learning introduces new challenges, including the ability of the model to memorize some specific training data points and thus be vulnerable to privacy attacks that seek to extract this data from the model. Fortunately, techniques such as differential privacy exist that can be helpful in overcoming this specific challenge. Just as with other areas of Responsible AI, it’s important to be aware of these new challenges that come along with machine learning and what steps can be taken to mitigate them.

Credits

Adam Pearce and Ellen Jiang // December 2020

Thanks to Andreas Terzis, Ben Wedin, Carey Radebaugh, David Weinberger, Emily Reif, Fernanda Viégas, Hal Abelson, Kristen Olson, Martin Wattenberg, Michael Terry, Miguel Guevara, Thomas Steinke, Yannick Assogba, Zan Armstrong and our other colleagues at Google for their help with this piece.

More Explorables

Can a Model Be Differentially Private and Fair?

Training models with differential privacy stops models from inadvertently leaking sensitive data, but there's an unexpected side-effect: reduced accuracy on underrepresented subgroups.

Imagine you want to use machine learning to suggest new bands to listen to. You could do this by having lots of people list their favorite bands and using them to train a model. The trained model might be quite useful and fun, but if someone pokes and prods at the model in just the right way, they could extract the music preferences of someone whose data was used to train the model. Other kinds of models are potentially vulnerable; credit card numbers have been pulled out of language models and actual faces reconstructed from image models.

Training with differential privacy limits the information about any one data point that is extractable but in some cases there’s an unexpected side-effect: reduced accuracy with underrepresented subgroups disparately impacted.

Recall that machine learning models are typically trained with gradient descent, a series of small steps taken to minimize an error function. To show how a model can leak its training data, we’ve trained two simple models to separate red and blue dots using two simple datasets that differ in one way: a single isolated data point in the upper left has been switched from red to blue.

Notice that the two models have very different boundary lines near the isolated point by the end of the training. Someone with access to the trained model might be able to infer if the point in the upper left is red or blue — if the color represented sensitive information, like someone’s voting record, that could be quite bad!

Protecting the Privacy of Training Points

We can prevent a single data point from drastically altering the model by adding two operations to each training step:²

⚬ Clipping the gradient (here, limiting how much the boundary line can move with each step) to bound the maximum impact a single data point can have on the final model.
⚬ Adding random noise to the gradient.

Try increasing the random noise below. We’re now training lots of differentially private models; the more the potential models for the red and blue outlier points overlap, the more plausible deniability the person in the upper left has.

You can also try dragging the other points around and adjusting the gradient clipping. Are points in the center or outliers more likely to modify the boundary lines? In two dimensions there’s a limited number of outliers, but in higher dimensions more points are outliers and much more information can be extracted from a trained model.

Correctly combined, adding gradient clipping and random noise to gradient descent make it possible to train a model with differential privacy – we can guarantee that a model trained on a given dataset is essentially indistinguishable from a model trained on the same dataset with a single point changed.

Predictions on Outliers Change the Most

What does this look like in practice? In Distribution Density, Tails, and Outliers in Machine Learning, a series of increasingly differentially private models were trained on MNIST digits. Every digit in the training set was ranked according to the highest level of privacy that correctly classified it.

On the lower left, you can see digits labeled as “3” in the training data that look more like a “2” and a “9”. They’re very different from the other “3”s in the training data so adding just a bit of privacy protection causes the model to no longer classify them as “3”. Under some specific circumstances, differential privacy can actually improve how well the model generalizes to data it wasn’t trained on by limiting the influence of spurious examples.

The right side shows more canonical digits which are classified correctly even with high levels of privacy because they’re quite similar to other digits in the training data.

The Accuracy Tradeoff

Limiting how much a model can learn from a single example does have a downside: it can also decrease the model’s accuracy. With 7,500 training points, 90% accuracy on MNIST digits is only achievable with an extremely low level of privacy protection; increasing privacy quickly lowers the model’s accuracy.

Collecting more training data offers a way out of this accuracy/privacy tradeoff. With 60,000 training points, 90% accuracy can be reached with a higher privacy level than almost all real-world deployments of differential privacy.

Looking at the differences between predictions by digit class shows another potential complication: some classes are harder to identify than others. Detecting an “8” with high confidence requires more training data and/or lower privacy than detecting a “0” with high confidence.

This problem is exacerbated if the training data has fewer examples of one class than the others. Trying to predict an uncommon event with a differentially private model can require an enormous amount of data.

Implications for Fairness

Outliers also aren’t evenly distributed within a class. Below, MNIST digits are colored by their sensitivity to higher privacy levels and projected with UMAP, forming several clusters of privacy-sensitive yellow digits. It’s possible to inadvertently train a model with good overall accuracy on a class but very low accuracy on a smaller group within the class.

There’s nothing that makes a “1” slanted to the left intrinsically harder to classify, but because there are only a few slanted “1”s in the training data it’s difficult to make a model that classifies them accurately without leaking information.

This disparate impact doesn’t just happen in datasets of differently drawn digits: increased levels of differential privacy in a range of image and language models disproportionality decreased accuracy on underrepresented subgroups. And adding differential privacy to a medical model reduced the influence of Black patients’ data on the model while increasing the influence of white patients’ data.

Lowering the privacy level might not help non-majoritarian data points either – they’re the ones most susceptible to having their information exposed. Again, escaping the accuracy/privacy tradeoff requires collecting more data – this time from underrepresented subgroups.

Credits

Adam Pearce // January 2022

Thanks to Abhradeep Thakurta, Andreas Terzis, Andy Coenen, Asma Ghandeharioun, Brendan McMahan, Ellen Jiang, Emily Reif, Fernanda Viégas, James Wexler, Kevin Robinson, Matthew Jagielski, Martin Wattenberg, Meredith Morris, Miguel Guevara, Nicolas Papernot and Nithum Thain for their help with this piece.

Footnotes

To speed up training at the cost of looser privacy bounds, gradients, clipping and noise can be calculated on a group of data points instead of individual data points.

The “ε” in ε-differential privacy essentially measures the overlap in two distributions after changing a single data point.

Clipping and noising are also used outside of differential privacy as regularization techniques to improve accuracy.

In addition to accidently mislabeled examples, differential privacy can also provide some protection against data poisoning attacks.

While visually similar digits aren’t necessarily interpreted in similar ways by the model, the clustering of visually similar digits in the UMAP diagram at the bottom of the page (which projects embedding from the penultimate layer of digit classifier) suggests there is a close connection here.

Rebalancing the dataset without collecting more data doesn’t avoid this privacy/accuracy tradeoff – upsampling the smaller class reduces privacy and downsampling the larger class reduces data and lowers accuracy.

See the appendix on Subgroup Size and Accuracy for more detail.

Appendix: Subgroup Size and Accuracy

How, exactly, does the amount of training data, the privacy level and the percentage of data from a subgroup impact accuracy? Using MNIST digits rotated 90° as a stand-in for a smaller subgroup, we can see how the accuracy of a series of simple models that classify “1”s and “7”s change based on these attributes.

On the far left, models without any rotated digits in the training data never classify those digits more accurately than random guessing. By rotating 5% of the training digits, a small slice of models with lots of training data and low privacy can accurately classify rotated digits.

Increasing the proportion of rotated digits to 10% or 20% or even more makes it possible to train a higher privacy model that performs well on both types of digits with the same amount of training data.

Click on one of the models above and you can see how the accuracy gap shifts as number of training points, privacy level and percentage of rotated digits are independently changed.

Intuitively, adding more training data has diminishing marginal increases to accuracy. Accuracy on the smaller group of rotated digits, which may just be on the cusp of being learned, falls off faster as the effective amount of training data is decreased — a disparate reduction in accuracy.

More Explorables

Federated Learning

How Federated Learning Protects Privacy

Large datasets have made astounding breakthroughs in machine learning possible. But oftentimes data is personal or proprietary, and not meant to be shared, making privacy a critical concern of and barrier to centralized data collection and model training. With federated learning, it’s possible to collaboratively train a model with data from multiple users without any raw data leaving their devices. If we can learn from data across many sources without needing to own or collect it, imagine what opportunities that opens!

Billions of connected devices — like phones, watches, vehicles, cameras, thermostats, solar panels, telescopes — with sensors to capture data and computational power to participate in training, could collaborate to better understand our environment and ourselves. How do people move? What impacts our health and wellbeing? Together via federated learning these devices could enable new technologies as well. Consider how our cars might contribute to large-scale training of autonomous vehicles without divulging our whereabouts.

And this machine learning approach can be applied across separate organizations as well. Hospitals could design better treatment plans with knowledge about patient outcomes due to various interventions from care providers worldwide without sharing highly sensitive health data. Pharmaceutical companies with proprietary drug development data could collaborate to build knowledge about how the body is likely to metabolize different compounds.

This framework has the potential to enable large-scale aggregation and modeling of complicated systems and processes like urban mobility, economic markets, energy use and generation patterns, climate change and public health concerns. Ultimately, the hope of federated learning is to allow people, companies, jurisdictions and institutions to collaboratively ask and answer big questions, while maintaining ownership of their personal data.

Designing a Federated Learning System

Let’s explore how this technology works with a simple example we can all relate to: blocking spam messages. Spam in chat apps is annoying and pervasive. Machine learning offers a solution – we could develop a model that automatically filters out incoming spam based upon what users previously marked as spam on their devices. This sounds great, but there’s a catch: most machine learning models are trained by collecting vast amounts of data on a central server; and user messages can be quite personal. To protect privacy, is it possible to train a spam detection model — or any machine learning model, for that matter — without sharing any potentially sensitive information with a central server?

To answer this question, let’s first take a closer look at a typical centralized training system, illustrated by the simple spam detection model below. User messages are uploaded to a central server, where they’re processed all at once to train a bag-of-words model. Click a message to flag it as spam ❌ or not to change the data uploaded to the server and the trained model.

This model might be pretty good at filtering out spam messages. But centralized training comes with a big downside: all the messages, no matter how sensitive, need to be sent to the server, requiring users to trust the owners of that centralized server to protect their data and not misuse it.

What if training was done locally on each user’s device instead and their data wasn’t centrally collected? Smartphones are getting increasingly powerful, and they’re often idle — for example, while charging overnight — enabling machine learning model training to run without impacting the user experience.

Training models locally is great for privacy — no data ever leaves a user’s device! — but we can see here how a single device with limited data might not be able to train a high quality model. If a new scam involving, say, car insurance starts spamming messages to everyone, Alice’s phone wouldn’t be able to filter out messages about “your auto warranty renewal” with a local-only model until she marks several of them as spam — even if Bob has already flagged similar messages.

How can users help each other out and collaboratively train a model without sharing their private data? One idea is for users to share their locally trained spam-detection models instead of their messages. The server can then combine these models, for example by averaging them, to produce a global model that everyone could use for spam filtering.

While we’ve stopped sending every raw message to the server, uploading these local models still leaks some information. Here, the central server has direct access to the rates each user marks different words as spam and can infer what they’re talking about. Depending on the level of trust users have in the server, they may be uncomfortable with the server seeing their local models. Ideally the server should only see the aggregated result. We want to develop a system that provides as much data minimization as possible.

Federated learning is a general framework that leverages data minimization tactics to enable multiple entities to collaborate in solving a machine learning problem. Each entity keeps their raw data local, and improves a global model with focused updates intended for immediate aggregation. A good first step towards limiting data exposure when combining user models is to do so without ever storing the individual models — only the aggregate. Secure aggregation and secure enclaves can provide even stronger guarantees, combining many local models into an aggregate without revealing the contribution of any user to the server. This may sound almost magical, so let’s take a closer look at how secure aggregation works.

In the secure aggregation protocol, user devices agree on shared random numbers, teaming up to mask their local models in a way that preserves the aggregated result. The server won’t know how each user modified their model.

Try dragging the shared random numbers — the aggregated sum remains constant even though what each user sends to the server changes. And importantly, the users’ hidden numbers are never shared!

Let’s put everything together by running all the user-contributed numbers that make up each of their local models through the secure aggregation process. Alice, Bob and Carol’s devices use a cryptographic technique to exchange random numbers secretly — our users won’t actually meet in person.

With secure aggregation, users collaboratively merge their models without revealing any individual contribution to the central server.

All in all, federated learning enables collaborative model training, while minimizing data exposure. This system design extends to problem settings far beyond the toy spam example we’ve illustrated above, to larger-scale modeling across all sorts of devices and institutions with privately-held data.

Training a Federated Model

While a very simple model like our toy spam classifier can be learned via a single round of merging local models, more sophisticated models require many iterations of local training and federated averaging. Let’s see how that works and examine some challenges that arise in practice. We’ll look at a simple “heat-map” binary classification model designed to guess what regions of a grid are likely to be hot or cold. Each of our users has only collected temperature readings from a handful of locations:

If all the users uploaded their data to a central server, it’d be easy to spot the pattern:

Our goal is to learn this temperature distribution across the grid — so everyone will know where they need a sweater! — without anyone having to share their location history.

Below, each user is continually training a model with just their local data, predicting the temperature of every location in the grid. You can see how dramatically different models are trained as each user’s model overfits to their limited view of the world. The local training curves track the accuracy of each local model on the ground truth data, indicating how well each local model learns the true temperature distribution across the grid.

Click to run a round of federated training: averaging user models and distributing the updated global model to all users. After training and merging models several times, the resulting global model better resembles the overall temperature distribution across the map than the models trained on just local data. You may notice how local heat map models drift apart after a significant period of local training, and the latest global model’s accuracy might degrade upon merging. Relatively frequent periodic averaging is used to avoid this.

While we plot the local model accuracies so it’s possible to observe these training dynamics, in practice a server running federated training only has access to the global model. The only metric that can be computed and tracked over the course of training by the server is the global model accuracy.

Outliers

This works pretty well when all users report consistent temperature experiences. What happens if that’s not the case? Maybe some of our users have broken thermometers and report cold weather everywhere! Click on each of the four outliers to exclude them from training and notice how the model performs.

We may be able to better train a model to predict the heat-map that the majority of users observed without the outliers, but what if these outlier users don’t have broken sensors and their data just looks different? Some people may have different ideas of what is “hot” or “cold;” excluding outliers from training risks reducing accuracy for groups of people less represented in the training pool.

Though it’s easy to spot the outliers in this example, in practice the server in a federated learning system cannot directly see user training data, which makes detecting outliers in federated learning tricky. The presence of outliers is often indicated by poor model quality across users.

Differential Privacy

Having the global model drastically change based on the presence of a single user also raises privacy concerns. If one user’s participation can significantly affect the model, then someone observing the final model might be able to determine who participated in training, or even infer their local data. Outlier data is particularly likely to have a larger impact on model training.

For example, let’s say our potential user group includes one person known to always wear a sweater and complain about the cold. If the global model accuracy is lower than expected, we can infer that the notorious sweater-wearing user probably participated in training and reduced accuracy by always reporting cold. This is the case even with secure aggregation — the central server can’t directly see which user contributed what, but the resulting global model still gives away that it’s likely that a user who believes that it’s always sweater weather participated.

Carefully bounding the impact of any possible user contribution and adding random noise to our system can help prevent this, making our training procedure differentially private. When using differential privacy in federated learning, the overall accuracy of the global model may degrade, but the outcome should remain roughly the same when toggling inclusion of the outlier (or any other user) in the training process.

Use the slider to modulate how much the user-reported locations are perturbed. At lower levels of privacy toggling the inclusion of the outlier affects the model more significantly, whereas at higher levels of privacy there is not a discernible difference in model quality when the outlier is included.

In practice user models are clipped and noised rather than their raw data, or noise is applied to the combination of many clipped models. Applying the noise centrally tends to be better for model accuracy, however the un-noised models may need to be protected by technologies like trusted aggregators.

This demonstration illustrates a trade-off between privacy and accuracy, though there’s another missing dimension that factors into the equation: the amount of data, both in number of training examples and number of users. The cost of using more data isn’t free — this increases the amount of compute — but it’s another knob we can turn to arrive at an acceptable operating point across all of these dimensions.

Playground

There are lots of other knobs to turn in a federated learning setting. All these variables interact in complicated ways. Click on a value for each variable to run a particular configuration, or the variable name to sweep over all of its options. Go ahead and play around — try mixing them together!5

Credits

Nicole Mitchell and Adam Pearce // November 2022

Thanks to Nithum Thain, Alex Ingerman, Brendan McMahan, Hugo Song, Daniel Ramage, Peter Kairouz, Alison Lentz, Kallista Bonawitz, Jakub Konečný, Zachary Charles, Marco Zamarato, Zachary Garrett, Lucas Dixon, James Wexler, Martin Wattenberg, Astrid Bertrand and the Quirk Research team for their help with this piece.

Footnotes

For an end-to-end encrypted messaging application, centralized learning won’t work.

More cryptographically sophisticated protocols can be used so a connection between every user isn’t required and the sums can still be computed if some users drop out.

Curious how these models are being trained and what’s going on in each “local step”? Check out the TensorFlow Playground.

In the cross-device setting (e.g., billions of smartphones), a small fraction of all devices are sampled to participate in each round; typically any one device will contribute to training a handful of times at most. This spreads out the load of training, and ensures the model sees a diversity of different devices. In the cross-silo setting, a small number of larger and more reliable users are assumed (e.g., organizations, datacenters).

In a real cross-device FL training system, each sampled device would generally only compute a fixed relatively small number of local steps before averaging.

The system below merges the local models every 20 local training steps.

Simply including all types of users is often not enough to ensure fairness. Designing strategies for learning a model that performs equally well for everyone is an active area of research. Personalization through local fine-tuning of a final global model is one promising approach.

This might not seem like the biggest deal for this scenario, but participation in medical trials can be highly sensitive and more complex models can leak information.

It’s important to remember that this is a simple model and a very small scale federated learning simulation. The phenomena you observe might not be fully representative of what happens in practice.

More Explorables

Are Model Predictions Probabilities?

If a machine learning model tells you that it’s going to rain tomorrow with a score of 0.60, should you buy an umbrella?1

In the diagram, we have a hypothetical machine learning classifier for predicting rainy days. For each date, the classifier reads in relevant signals like temperature and humidity and spits out a number between 0 and 1. Each data point represents a different day, with the position representing the model’s prediction for rain that day and the symbol (🌧️ or ☀️) representing the true weather that occurred that day.

Do the model’s predictions tell us the probability of rain?

In general, machine learning classifiers don’t just give binary predictions, but instead provide some numerical value between 0 and 1 for their predictions. This number, sometimes called the model score or confidence, is a way for the model to express their certainty about what class the input data belongs to. In most applications, the exact score is ignored and we use a threshold to round the score to a binary answer, yes or no, rain or not. However, by using calibration we can transform these scores into probabilities and use them more effectively in decision making.

Thresholding

One traditional approach to using a model’s score is through thresholding. In this setting, you choose a threshold t and then declare that the model thinks it’s going to rain if the score is above t and it’s not if the score is below, thereby converting the score to a binary outcome. When you observe the actual weather, you know how often it was wrong and can compute key aggregate statistics like accuracy.

We can sometimes treat these aggregate statistics themselves as probabilities. For example, accuracy is the probability that the binary prediction of your model (rain or not) is equal to the ground truth (🌧️ or ☀️).

Adjustable Thresholding

The threshold can easily be changed after the model is trained.

Thresholding uses the model’s score to make a decision, but fails to consider the model’s confidence. The model score is only used to decide whether you are above or below the threshold, but the magnitude of the difference isn’t considered. For example, if you threshold at 0.4, the model’s predictions of 0.6 and 0.9 are treated the same, even though the model is much more confident in the latter.

Can we do a better job of incorporating the model score into our understanding of the model?

Calibration

Calibration lets us compare our model scores directly to probabilities.

For this technique, instead of one threshold, we have many, which we use to split the predictions into buckets. Again, once we observe the ground truth, we can see what proportion of the predictions in each bucket were rainy days (🌧️). This proportion is the empirical probability of rain for that bucket.

Ideally, we want this proportion to be higher for higher buckets, so that the probability is roughly in line with the average prediction for that bucket. We call the difference between the proportion and the predicted rates the calibration error, and by averaging over all of the buckets, we can calculate the Expected Calibration Error. If the proportions and the predictions line up for our use case, meaning the error is low, then we say the model is “well-calibrated” and we can consider treating the model score as the probability that it will actually rain.

Adjusting Calibration

We saw above that a well-calibrated model allows us to treat our model score as a kind of probability. But if we start with a poorly calibrated model, one which is over or under-confident. Is there anything we can do to improve it?

It turns out that, in many settings, we can adjust the model score without really changing the model’s decisions, as long as our adjustment preserves the order of the scores2. For example, if we map all of the scores from our original model to their squares, we don’t change the order of the data with respect to the model score. Thus, quantities like accuracy will stay the same as long as we appropriately map the threshold to its square as well. However, these adjustments do change the calibration of a model by changing which data points lie in which buckets.

Try tweaking the thresholds to calibrate the model scores for our data3 – how much can you improve the model’s calibration?

In general, we don’t have to rely on tweaking the model scores by hand to improve calibration. If we are trying to calibrate the model for a particular data distribution, we can use mathematical techniques like Isotonic Regression or Platt Scaling to generate the correct remapping for model scores.

Shifting Data

While good calibration is an important property for a model’s scores to be interpreted as probabilities, it alone does not capture all aspects of model uncertainty.

What happens if it starts to rain less frequently after we’ve trained and calibrated our model? Notice how the calibration drops, even if we use the same calibrated model scores as before.

Models are usually only well calibrated with respect to certain data distributions. If the data changes significantly between training and serving time, our models might cease to be well calibrated and we can’t rely on using our model scores as probabilities.

Beyond Calibration

Calibration can sometimes be easy to game. For example, if we knew that it rains 50% of the time over the course of the year, then we could create a model with a constant prediction of 0.5 every day. This would have perfect calibration, despite not being a very useful model for distinguishing day-to-day differences in the probability of rain. This highlights an important issue:

Better calibration doesn’t mean more accurate predictions.

It turns out that statisticians identified the issue with focusing solely on calibration in meteorology when comparing weather forecasts, and came up with a solution. Proper scoring rules provide an alternative approach to measuring the quality of probabilistic forecasts, by using a formula to measure the distance between the model’s predictions and the true event probabilities. These rules guarantee that a better value must mean a better prediction in terms of accuracy and calibration. Such rules incentivize models to be both better calibrated and more accurate.

Credits

Nithum Thain, Adam Pearce, Jasper Snoek & Mahima Pushkarna // March 2022

Thanks to Balaji Lakshminarayanan, Emily Reif, Lucas Dixon, Martin Wattenberg, Fernanda Viégas, Ian Kivlichan, Nicole Mitchell, and Meredith Morris for their help with this piece.

Footnotes

Your decision might depend both on the probability of rain and its severity (i.e. how much rain there is going to be). We’ll focus just on the probability for now.

Applying a strictly monotonic function to the model always keeps the order of scores the same.

In this example, we adjust the model scores by changing the model scores of elements within a bucket to the mean of the bucket.

More Explorables

Searching for Unintended Biases With Saliency

Machine learning models are used for high stakes tasks like steering driverless cars or detecting cancerous tissue from medical scans. If there are spurious correlations1 in the training data, the model might develop unintended biases that could lead to mistakes. In this post, we explore a technique for detecting these biases by asking for an explanation of how models make decisions.

We will explore this phenomenon through a pair of simple classification models which are trained to detect if an image is of a cat 🐱 or a dog 🐶.

Below are both models’ predictions2 on a test set of 50 images. Can you tell which model would be better to deploy to users?

There’s not much of a difference; it’s hard to say which model would perform better in a real-world setting. To differentiate these two models, let’s use a different collection of images to check whether the predictions have an unintended bias from spurious correlations in their training datasets.

Now in almost every case, Model B predicts that the image is a dog!

Looking closely at the first set of images, there’s a small but significant difference between the cat and dog images.

All the cat images have a watermark in the corner, and it seems like Model B learned to detect watermarks instead of cats. As a matter of fact, Model B, which we’ll call the Watermark Model, was trained on a dataset where cats had watermarks and picked up a bias relating watermarks to cats. Model A correctly recognizes cats in images without watermarks and doesn’t appear to use any shortcut. This model, which we’ll call the Normal Model, was trained on watermark-free images so it didn’t pick up the watermark shortcut. It would be a much better choice to deploy.

In this contrived example, we were lucky that we had access to an unbiased dataset. It enabled us to form a preliminary hypothesis to explain why Model B performs so poorly. But what if we didn’t have an unbiased dataset?

Next, we’ll look at a set of tools that will show us which areas of an image a model relies on. If we see something like a watermark being used to differentiate between cats and dogs, we’ll get some insights into potential problems with our model.

Occlusion

How can we tell which parts of an image a model is using? One simple method is to occlude part of the image with a black box and check how the model’s prediction changes.

Try mousing over different parts of this cat — what do you need to hide to make the Watermark Model predict dog? What about for the Normal Model?

The Watermark Model prediction changes if the bottom left corner is covered up, which provides evidence that the model is relying on a spurious correlation.

Manually checking each box is slow, however. We can speed things up by automatically checking the boxes one by one and overlaying the results on the image instead.

While occlusion gives us some preliminary explanation of the model’s reasoning, it’s far from perfect. At higher resolutions, it doesn’t highlight the lower left for the Watermark Model. This is because occlusion only shows the effect of covering up a single box and doesn’t account for what would happen if multiple boxes were occluded simultaneously.

Occlusion-based methods also require significant computational power. The model needs to be rerun each time we hide a cell in the grid, which gets time-consuming as the resolution of the grid increases.

Leveraging Gradients

Machine learning researchers have developed various techniques to visualize model decision making. One set of approaches tries to determine areas of an image that are most “salient” from a model’s perspective and produce a saliency map that’s similar to the overlaid occlusion predictions. They’re usually based on taking the gradient of the model’s prediction against an image.

The gradient gives us information about how the prediction would change if a tiny positive change is applied to an individual input image feature (i.e. the RGB values of a pixel), and it does this simultaneously for all image features. This simultaneity makes saliency maps much less computationally intensive than the occlusion method described above.

To better understand gradients, the image below demonstrates how the model’s prediction that the image below is a “cat” would change as we change each pixel. If a change is more likely to make the model think the image is a cat, we use an upward arrow. The thicker the arrow, the bigger a change in gradient for changing that particular pixel. We call pixels which make a big difference “salient”.

The simple gradient-based method is often referred to as Vanilla Gradient since saliency is determined solely by model gradients. Follow up work has built on this approach by transforming the gradients in different ways to improve the explanatory power of the saliency maps.

In our subsequent visualizations, we compute saliency maps over pixels, with white indicating salient pixels (those which most affect the prediction), and black indicating non-salience. Hover over the thumbnail images below to see what their Vanilla Gradient saliency maps look like for our two models.

As you can see, the Vanilla Gradient method tends to be a bit noisy.

There are some fairly simple transformations we can do over the Vanilla Gradient’s approach to reduce its noise. A simple one consists of taking the square of the gradient,3 which emphasizes higher values and focuses on the size of the gradient, ignoring the direction.4

For the Normal Model, saliency maps highlight various features of the images, possibly relevant to the prediction of cats and dogs (like eyes, nose, and body shape), but are quite hard to interpret.

Let’s take a closer look at what the saliency maps of the Watermark Model detect for images of cats 🐱 and dogs 🐶 with and without watermarks:

To faithfully reflect how the Watermark Model is making decisions, you might expect that saliency maps would highlight the watermark area even in watermark-free images. This is generally the case, although saliency maps are less precise and noisier for watermark-free images, especially with Vanilla Gradient.

More Subtle Bias

What if the bias were less obvious? Would simple pixel-based saliency maps still pick it up?

The Watermark Model was trained on a biased dataset where all cat images were watermarked. However, bias is usually more subtle. It rarely affects 100% of your training set but often appears more sporadically.

Below, we have a model trained on a dataset where 50% of the cat images are watermarked.

It appears that the model sometimes uses watermarks for its predictions, as we can see from the mistakes it makes, but how well is this detected by our saliency maps?

We can quantify the effectiveness of saliency maps in flagging watermarks with a simple metric — the proportion of “salient” pixels that are located within the watermark area.5 The recipe for this approach is: (1) we take the smallest rectangle including our suspected spurious correlation6 (the watermark) (2) we select all high gradient values from the saliency map (the “very white” pixels)7 (3) we count how many of these values are in our watermark rectangle.

We implemented this approach in the diagram below for the four different categories of test data. Each circle represents the model prediction on a different image belonging to that category. Hover over the dots to display its corresponding image and saliency map.

The 0% model, where none of the cat images in the training data have watermarks, is our Normal Model. All images of cats (left quadrants) are correctly classified as cats. Nearly all dog images (right quadrants) are correctly classified as dogs.

The 100% model, where all of the cat images in our training data have watermarks, is our Watermark Model. Unsurprisingly, we see that the model makes many mistakes, with the bottom left and top right quadrants misclassified. The points on the charts move to the right end of the x-axis, indicating that most of the high salience pixels lie in the watermark area of the image, helping us understand why our model makes these mistakes. However, the watermark is highlighted a little less clearly for images without a watermark. A few saliency maps in the bottom left quadrant don’t pick up at all the watermark, although images were misclassified.

When the bias is more sporadic in the training set (e.g. affecting only 50% of cat images), it gets even more difficult to detect the spurious correlation using saliency maps, especially when looking at watermark-free images. On the top right quadrant, most of the salient features for misclassified dogs are in the watermark region. However, on the bottom left quadrant, many cat images are misclassified as dogs, likely because they don’t have watermarks, but very few saliency maps actually highlight the watermark area. That’s one limitation of saliency maps: they are not very good at highlighting what’s missing.

Other Forms of Bias

Spurious correlations can take many forms, and are generally more subtle than watermarks. Can simple saliency maps alert us to other forms of bias?

In the diagram below, we have saliency maps for three “mystery models” on four input images.

Can you recognize if any of the models rely on spurious correlation? If so, which ones? When you’ve made up your mind, click on the column title to reveal the model’s characteristics.

You’ve probably noticed how challenging it is to detect biases from saliency maps alone.

For example, “Mystery Model 2” relies on the color of the animal to make its prediction, but the saliency maps for this model seem to highlight the animal’s face and body, which is probably consistent with what a human would consider meaningful. In this case, the saliency maps might even do more harm than good: they might have misled you into thinking that this model was making correct decisions based on pertinent features.

Even when saliency maps can correctly indicate spurious signals, it can be difficult to see those signals when you don’t know what you’re looking for. In several controlled experiments, Adebayo et al. found that saliency maps8 were unable to help people detect certain unknown spurious correlations.

Understanding why your model makes a decision is important for trusting your model, but how much you can trust the explanations themselves is also an important question, and an open research topic. We’ve seen in this post that saliency maps can be useful to detect some biases in datasets but also that it can be difficult to see bias when it’s more subtle and sporadic in your training set. While they can give you insights into what features a model is using or misusing to make its decisions, sometimes saliency maps simply don’t help you draw any conclusions about a model.

In general, it’s always helpful to thoroughly understand your training data. Tools like Know Your Data and LIT help researchers, engineers, product teams and policy teams explore datasets and model predictions to improve data quality and mitigate bias issues.

Additionally, supplementing your analysis with several types of interpretability methods improves the likelihood of detecting errors. In the section below we provide several examples of other useful interpretability methods.

Beyond Simple Saliency Methods

Beyond the simple techniques presented in this post, a myriad of other saliency methods exist. They are broadly divided into three categories. Sensitivity methods, like Vanilla Gradient, show how a small change to the input affects the prediction. Signal methods, like DeConvNet or Guided BackProp, look at the neuron activations in the model to attribute the importance of input features. Finally, attribution methods, like Integrated Gradients 9 and SHAP aim at completely specifying the attributions for all the input features so that they sum up to the output.

Saliency methods can be applied to other types of data, like text. There’s also research focused on making saliency maps more “human-interpretable.” Looking at individual pixels is hard for people and difficult to interpret, so techniques like XRAI and LIME instead create maps that highlight the most important regions in the image.

While interpretability research is constantly producing new methods, a complementary line of work is dedicated to critically examining and measuring their limitations. Sanity Checks for Saliency Maps presents different experiments on saliency maps to check that they behave in the way we expect them to.

Furthermore, the research space in interpretability isn’t restricted to saliency methods. For example, influence methods, also known as training data attribution, suggest which training data points might be the cause of a model’s behavior for a given input and output. Some state-of-the-art examples of influence methods are this paper, this one or this one.

Researchers have also explored mapping models’ internal representation to human concepts. In the natural language domain, Bolukbasi et al. used relations between concepts to reduce bias in word embeddings. More recently, Kim et al. popularized the use of human-specified labels for image models, enabling the creation of classifiers for high level concepts like “whisker” or “paw.”

Credit

Astrid Bertrand, Adam Pearce and Nithum Thain // December 2022

Thanks to Ben Wedin, Tolga Bolukbasi, Nicole Mitchell, Lucas Dixon, Andrei Kapishnikov, Blair Bilodeau, Been Kim, Jasmijn Bastings, Katja Filippova and Seyed Kamyar Seyed Ghasemipour for their help with this piece.

Please cite as:

Astrid Bertrand, Adam Pearce and Nithum Thain. “Searching for Unintended Biases with Saliency” PAIR Explorables, 2022.

BibTeX:

@article{bertrand2022saliency,
   title={Searching for Unintended Biases with Saliency},
   author={Bertrand, Astrid and Pearce, Adam and Thain, Nithum},
   year={2022},
   journal={PAIR Explorables},
   note={https://pair.withgoogle.com/explorables/saliency/}
}

Images from Pexels and Kaggle.

Footnotes

“Spurious correlation” is a term to indicate when two variables are correlated but don’t have a causal relationship. In our case, watermarks and cats are spuriously correlated.

The confidence score is also shown.

Taking the square of the vanilla gradient produces less noisy images.

Other, more sophisticated methods exist to “denoise” Vanilla Gradient. For example, SmoothGrad (Smilkov et al., 2017) reduces variance through imperfect copies. The technique consists of taking the saliency maps of several copies of the input image where some noise was added, and averaging them together.

In this diagram we visualize the saliency maps for this model using the Gradient Squared method.

There are other measures we can use to evaluate saliency maps. “Known Spurious Signal Detection Measure” (K-SSD) is very similar: it measures the similarity of saliency maps derived from spurious models to an image where strictly the spurious signal is highlighted. “False Alarm Measure (FAM)” measures the similarity of explanations derived from normal models for spurious inputs to explanations derived from spurious models for the same inputs. See Adebayo et al, 2022 for the full definitions and Denain et al., 2022 for an implementation of the similarity measure using embeddings of saliency maps in a semantic feature space.

We take the 0.5% highest gradient values. The reason we take so little is that most of the gradient values are very close to 0 ( displayed in black in the saliency map). Only a very small fraction (0.5%) is closer to the maximum value.

When there are multiple sources of truth for the model, as in our 50% case, where the model uses both animal and watermark features, the model may only need one type of feature to make the prediction. This means that sometimes, it may not pay attention to the watermark, but actually consider the other important features it learned during training. Therefore, for watermark-free images the model may really infer that a cat is a dog not because there was no watermark but because of poor training.

They used the following saliency methods: Input-Gradient, SmoothGrad, Integrated Gradients (IG), and Guided Backprop (GBP).

Take a look at this post, which describes the Integrated Gradient method in more detail.

References

Adebayo, Julius, Justin Gilmer, Michael Muelly, Ian Goodfellow, Moritz Hardt, and Been Kim. “Sanity Checks for Saliency Maps.” arXiv, November 6, 2020. https://doi.org/10.48550/arXiv.1810.03292. Adebayo, Julius, Michael Muelly, Harold Abelson, and Been Kim. “Post Hoc Explanations May Be Ineffective for Detecting Unknown Spurious Correlation,” 2022. https://openreview.net/forum?id=xNOVfCCvDpM. Akyürek, Ekin, Tolga Bolukbasi, Frederick Liu, Binbin Xiong, Ian Tenney, Jacob Andreas, and Kelvin Guu. “Towards Tracing Factual Knowledge in Language Models Back to the Training Data.” arXiv, October 25, 2022. http://arxiv.org/abs/2205.11482. Bastings, Jasmijn, Sebastian Ebert, Polina Zablotskaia, Anders Sandholm, and Katja Filippova. “‘Will You Find These Shortcuts?’ A Protocol for Evaluating the Faithfulness of Input Salience Methods for Text Classification.” Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, 2022. https://doi.org/10.48550/arXiv.2111.07367. Bolukbasi, Tolga, Kai-Wei Chang, James Zou, Venkatesh Saligrama, and Adam Kalai. “Man Is to Computer Programmer as Woman Is to Homemaker? Debiasing Word Embeddings.” arXiv, July 21, 2016. https://doi.org/10.48550/arXiv.1607.06520. Denain, Jean-Stanislas, and Jacob Steinhardt. “Auditing Visualizations: Transparency Methods Struggle to Detect Anomalous Behavior.” arXiv, June 27, 2022. http://arxiv.org/abs/2206.13498. Kapishnikov, Andrei, Tolga Bolukbasi, Fernanda Viégas, and Michael Terry. “XRAI: Better Attributions Through Regions.” arXiv, August 20, 2019. https://doi.org/10.48550/arXiv.1906.02825. Kim, Been, Martin Wattenberg, Justin Gilmer, Carrie Cai, James Wexler, Fernanda Viegas, and Rory Sayres. “Interpretability Beyond Feature Attribution: Quantitative Testing with Concept Activation Vectors (TCAV).” arXiv, June 7, 2018. https://doi.org/10.48550/arXiv.1711.11279. Koh, Pang Wei, and Percy Liang. “Understanding Black-Box Predictions via Influence Functions.” arXiv, December 29, 2020. https://doi.org/10.48550/arXiv.1703.04730. Lundberg, Scott, and Su-In Lee. “A Unified Approach to Interpreting Model Predictions.” arXiv, November 24, 2017. https://doi.org/10.48550/arXiv.1705.07874. Pruthi, Garima, Frederick Liu, Satyen Kale, and Mukund Sundararajan. “Estimating Training Data Influence by Tracing Gradient Descent.” In Advances in Neural Information Processing Systems, 33:19920–30. Curran Associates, Inc., 2020. https://proceedings.neurips.cc/paper/2020/hash/e6385d39ec9394f2f3a354d9d2b88eec-Abstract.html. Ribeiro, Marco Tulio, Sameer Singh, and Carlos Guestrin. “‘Why Should I Trust You?’: Explaining the Predictions of Any Classifier.” arXiv, August 9, 2016. http://arxiv.org/abs/1602.04938. Schioppa, Andrea, Polina Zablotskaia, David Vilar, and Artem Sokolov. “Scaling Up Influence Functions.” Proceedings of the AAAI Conference on Artificial Intelligence 36, no. 8 (June 28, 2022): 8179–86. https://doi.org/10.1609/aaai.v36i8.20791. Simonyan, Karen, Andrea Vedaldi, and Andrew Zisserman. “Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps.” arXiv, April 19, 2014. https://doi.org/10.48550/arXiv.1312.6034. Smilkov, Daniel, Nikhil Thorat, Been Kim, Fernanda Viégas, and Martin Wattenberg. “SmoothGrad: Removing Noise by Adding Noise.” arXiv, June 12, 2017. https://doi.org/10.48550/arXiv.1706.03825. Springenberg, Jost Tobias, Alexey Dosovitskiy, Thomas Brox, and Martin Riedmiller. “Striving for Simplicity: The All Convolutional Net.” arXiv, April 13, 2015. http://arxiv.org/abs/1412.6806. Sundararajan, Mukund, Ankur Taly, and Qiqi Yan. “Axiomatic Attribution for Deep Networks.” arXiv, June 12, 2017. https://doi.org/10.48550/arXiv.1703.01365. Zeiler, Matthew D., and Rob Fergus. “Visualizing and Understanding Convolutional Networks.” arXiv, November 28, 2013. https://doi.org/10.48550/arXiv.1311.2901. Zhou, Yilun, Serena Booth, Marco Tulio Ribeiro, and Julie Shah. “Do Feature Attribution Methods Correctly Attribute Features?” arXiv, December 15, 2021. http://arxiv.org/abs/2104.14403.

More Explorables

Datasets Have Worldviews

Every dataset communicates a different perspective. When you shift your perspective, your conclusions can shift, too.

Suppose you have a dataset of shapes. They can either be shaded or unshaded. They look something like this:

You built a supervised machine learning classifier that will automatically classify each shape as shaded or unshaded. You call it the “Is-Shaded Classifier”.

Click “Run Classifier” to see how your model performs.

It’s not perfect— some of the shapes are definitely misclassified. You want to improve your model!

To do so, you want to know more about the kinds of mistakes your model is making.

Thinking About Bias

In training, you only gave your model the raw image of each shape and one ground truth label: shaded and unshaded. But maybe something about your model—the distribution of the training data you used, the architecture you chose, or how you set your hyperparameters—resulted in your model performing better on some shapes than others.

In fact, you’ve seen a lot of papers and articles citing issues of biased model performance between circles, triangles, and rectangles in shape data. One paper finds that shape detection algorithms tend to do worse on triangles; another article says color accuracy is an issue with circles. So you wonder: are there biases in your model’s misclassifications?

You want to make sure that your model is performing equally well across circles, triangles, and rectangles, so you decide to do a fairness analysis.

There’s just one issue: you don’t have labels for which of your shapes are circles, triangles, or rectangles.

So, you decide to send your data to data labelers.

You receive feedback from your data labeling team that they’re not sure what to do with the shapes that aren’t exactly circles, triangles, or rectangles.

For the shapes that are unclear, you can have them use their best guess or simply label them as “other”. Then, you can finally do some fairness analysis!

Below is the interface they see:

These shapes should be labeled…

If you go back and change the labelers’ instructions, which shapes do you perform worst on? Where do you find bias?

You notice that your results hinge on how you choose to classify the shapes in our data.

Because ultimately, this isn’t a world of only circles, triangles, and rectangles!

Thinking About Classification

What could we find out about our classifier’s performance if we used different categories altogether?

All shapes are basically…

Everything else should be labeled…

With each of the different categories, which shapes do you perform worst on? Where do you find bias?

Each way of categorizing your shapes takes a different stance about what’s important . Each one makes some features more important than others, it make some distinctions visible and other distinctions invisible, and make some things easy to classify while others become outliers.

And each one tells you something different about what kind of bias your classifier has!

Grouping and Regrouping

Here’s another way to look at the same results. We can draw all the shapes that were correctly classified above the dashed line, and all the incorrectly classified shapes below it.

We’re still looking at the same model making the same classification on the same shapes, so the same shapes stay above and below the line. But each way of grouping the results distributes the errors differently— each way tells you something different.

Labels Tell Stories

The decisions you make about classification, however small…

All shapes are basically…

…begin to shape others’ decisions…

…they shape the analysis you can do…

…and they shape the kinds of conversations that happen.

It’s natural to want to find a way out of this problem by gathering more features or collecting more data. If we just have enough detail on enough data, surely we can avoid making these kinds of decisions, right?

Unfortunately, that isn’t the case. Describing the world around us in any way—whether we’re telling a friend a story or telling a computer about shapes—requires us to choose what information is important to convey and what tools we want to use to convey it.

Whether we think about it or not, we’re always making choices about classification.

And as we saw with shapes, all of these choices make some features more important than others, make some distinctions visible and other distinctions invisible, and make some things easy to classify while others become outliers.

In Practice

Let’s take a closer look at how this plays out in real machine learning applications. One straightforward example is in supervised object detection tasks.

For example, let’s imagine we want to train an object detection model on a dataset including this image:

Source: Wikimedia Commons

We could give it the following ground truth bounding boxes:

This looks objective, right? After all, a building is a building, a bush is a bush, and a mountain is a mountain!

But even labeling the same regions in the same image, you can communicate a very different perspective:

Or consider the image below, with several sets of “ground truth” labels. Looking at each of these labels, consider:

What features matter? What gets labeled? Whose worldview comes through? What might you learn from this set of labels that you wouldn’t learn from another?

Source: Wikimedia Commons

There is no “view from nowhere”, no universal way to organize every object, or word, or image. Datasets are always products of a particular time, place, and set of conditions; they are socially situated artifacts. They have histories; they have politics. And ignoring this fact has very real consequences.

So what do we do with this information?

A great place to start is to reflect on your own context and get curious about your data.

If it’s hard to see a dataset’s values—if it feels “objective”, “universal”, or “neutral”—it may simply be reflecting a worldview you’re accustomed to. So, understanding the limitations of your own worldview can tell you about the limitations of “objective” data. What assumptions do you make about the world? What feels like common sense? What feels foreign?

And do some sleuthing about your data! Who collected this data? Why was it collected? Who paid for it? Where did the “ground truth” come from?

You might even find yourself questioning what kinds of assumptions underpin machine learning dataset development or even thinking more deeply about classification as a whole.

If you find yourself with lots of questions, you’re already off to a good start.

Credits

Dylan Baker // January 2022

Thanks to Adam Pearce, Alex Hanna, Emily Denton, Fernanda Viégas, Kevin Robinson, Nithum Thain, Razvan Amironesei, and Vinodkumar Prabhakaran for their help with this piece.

More Explorables

Measuring Diversity

Search results that reflect historic inequities can amplify stereotypes and perpetuate under-representation. Carefully measuring diversity in data sets can help.

Search, ranking and recommendation systems can help find useful documents in large datasets. However, these datasets reflect the biases of the society in which they were created and the systems risk re-entrenching those biases. For example, if someone who is not a white man searches for “CEO pictures” and sees a page of white men, they may feel that only white men can be CEOs, further perpetuating lack of representation at companies’ executive levels.

Using the careful quantification outlined in a recent paper, Diversity and Inclusion Metrics in Subset Selection, we can quantify biases and push these systems to return a wider range of results.

The mathematics of all this is a little easier to follow with abstract shapes. Let’s take a look at some of them:

Suppose we want to return about 30% green boxes to reflect the distribution of some larger universe of shapes. Try clicking on the shapes below to select some of them — can you find a better subset to return?

Another diversity metric we care about is the percentage of dots… how close to 35% dots can you get?

If we can only return a single subset, how should we consider multiple diversity metrics? Sometimes it isn’t possible to reduce the difference of every metric to zero. One natural approach: find the selection with the lowest mean difference across all the metrics to get as close as possible to all the targets.

In other circumstances, like picking a panel of speakers, avoiding badly representing any single category might be more important. This can be done by finding the subset with the lowest max difference. Try minimizing both below:

Notice that minimizing the mean results in a different subset than minimizing the max; how else might using one over the other change the results?

Ranking Measures

We can pull out more detail by showing how the mean difference and maximum difference rank lots of sets. Below, there are 20 sets of 10 shapes sorted by the two measures. Try adjusting the target slider on the left to see how the rankings change; each set’s percentage of green, dots and small shapes are shown in the small histograms.

At the extremes, the choice of measure can have a big impact: if we want to try and return all green results, we can shift the green target up to 100%. With this target, the minimum difference basically sorts the sets by the number of green items and uses the other targets as a tiebreaker. In contrast, sorting by the mean difference balances the green target more with the dot and small targets.

Beyond mean and max differences, there are more ways to combine diversity metrics, like taking the cross of two metrics to account for intersectionality. The absolute value of the difference in target and actual percentages can also be quantified in other ways — you might want to penalize undershooting more than overshooting, for example. It’s important to keep in mind what exactly you’re trying to maximize and the dataset that you’re operating on.

Which Measure is Best?

In a vacuum, all of these ranking methods are defensible. Picking one requires knowledge of the dataset and broader societal context.

For example, the doctors on the left have more variance along the shirt color attribute, but they’re less diverse by gender than the doctors on the right. With the shirt color and gender targets we’ve picked, the two subsets have the same mean and max differences However, in most applications, it’s more important to have a representative sample of socially relevant characteristics, like gender, rather than something less salient, like clothing color.

Just selecting a diverse sample isn’t sufficient either. Diversity and Inclusion Metrics in Subset Selection introduces a way of measuring “inclusion” - how well does the searcher feel represented in the results?

Below, we have gender diversity, without inclusion for women, in the “construction worker” image domain. Masculine-presenting individuals are shown in realistic, modern construction worker situations, while feminine-presenting individuals and other gender presentations are depicted as historic nostalgia, toys, clipart, or passive.

The context of the query and the searcher also plays in the quality of search results. A search for “work clothing” that shows a mixed palette of colors for men’s clothing and only pink women’s clothing might make the searcher feel that women need to appear stereotypically feminine in a professional setting. But the same set of women’s clothes might be appropriate to show for a “pink women work clothes” search or if the searcher had previously expressed a preference for pink.

We saw how a small switch from mean to max made a huge difference in what abstract shapes are returned – and how things can get even more complex when socially salient characteristics are layered in. Defaults and small decisions can encode our priorities and values; intentionally thinking about how diversity and inclusion are being measured and which characteristics are emphasized is a step towards designing more equitable systems.

Credits

Adam Pearce, Dylan Baker, Ellen Jiang, Meg Mitchell* and Timnit Gebru* // March 2021

*Work done while at Google

Thanks to Alex Hanna, Carey Radebaugh, Emily Denton, Fernanda Viégas, James Wexler, Jess Holbrook, Ludovic Peran, Martin Wattenberg, Michael Terry, Yannick Assogba and Zan Armstrong for their help with this piece.

More Explorables

Measuring Fairness

There are multiple ways to measure accuracy. No matter how we build our model, accuracy across these measures will vary when applied to different groups of people.

Measuring Fairness

How do you make sure a model works equally well for different groups of people? It turns out that in many situations, this is harder than you might think.

The problem is that there are different ways to measure the accuracy of a model, and often it's mathematically impossible for them all to be equal across groups.

We'll illustrate how this happens by creating a (fake) medical model to screen these people for a disease.

Ground Truth

About half of these people actually have the disease a; half of them don't b.

Model Predictions

In a perfect world, only sick people would test positive for the disease and only healthy people would test negative.

Model Mistakes

But models and tests aren't perfect.

The model might make a mistake and mark a sick person as healthy c.

Or the opposite: marking a healthy person as sick f.

Never Miss the Disease...

If there's a simple follow-up test, we could have the model aggressively call close cases so it rarely misses the disease.

We can quantify this by measuring the percentage of sick people a who test positive g.

...Or Avoid Overcalling?

On the other hand, if there isn't a secondary test, or the treatment uses a drug with a limited supply, we might care more about the percentage of people with positive tests who are actually sick g .

These issues and trade-offs in model optimization aren't new, but they're brought into focus when we have the ability to fine-tune exactly how aggressively disease is diagnosed.

Try adjusting how aggressive the model is in diagnosing the disease

Subgroup Analysis

Things get even more complicated when we check if the model treats different groups fairly.¹

Whatever we decide on in terms of trade-offs between these metrics, we'd probably like them to be roughly even across different groups of people.

If we're trying to evenly allocate resources, having the model miss more cases in children than adults would be bad! ²

Base Rates

If you look carefully, you'll see that the disease is more prevalent in children. That is, the "base rate" of the disease is different across groups.

The fact that the base rates are different makes the situation surprisingly tricky. For one thing, even though the test catches the same percentage of sick adults and sick children, an adult who tests positive is less likely to have the disease than a child who tests positive.

Imbalanced Metrics

Why is there a disparity in diagnosing between children and adults? There is a higher proportion of well adults, so mistakes in the test will cause more well adults to be marked "positive" than well children (and similarly with mistaken negatives).

To fix this, we could have the model take age into account.

Try adjusting the slider to make the model grade adults less aggressively than children.

This allows us to align one metric. But now adults who have the disease are less likely to be diagnosed with it!

No matter how you move the sliders, you won't be able to make both metrics fair at once. It turns out this is inevitable any time the base rates are different, and the test isn't perfect.

There are multiple ways to define fairness mathematically. It usually isn't possible to satisfy all of them.³

Conclusion

Thankfully, the notion of fairness you choose to satisfy will depend on the context of your model, so while it may not be possible to satisfy every definition of fairness, you can focus on the notions of fairness that make sense for your use case.

Even if fairness along every dimension isn't possible, we shouldn't stop checking for bias. The Hidden Bias explorable outlines different ways human bias can feed into an ML model.

Appendix

¹ This essay uses very academic, mathematical standards for fairness that don't encompass everything we might include in the colloquial meaning of fairness. There's a gap between the technical descriptions of algorithms here and the social context that they're deployed in.

² Sometimes we might care more about different error modes in different populations. If treatment is riskier for children, we'd probably want the model to be less aggressive in diagnosing.

³The above example assumes the model sorts and scores people based on how likely it is that they are sick. With complete control over the model's exact rate of under- and over-diagnosing in both groups, it's actually possible to align both of the metrics we've discussed so far. Try tweaking the model below to get both of them to line up.

Adding a third metric, the percentage of well people a who test negative e, makes perfect fairness impossible. Can you see why all three metrics won't align unless the base rate of the disease is the same in both populations?

Drag ⁠— to adjust model accuracy and ⁠| to adjust the occurrence of disease

Credits

Adam Pearce // May 2020

Thanks to Carey Radebaugh, Dan Nanas, David Weinberger, Emily Denton, Emily Reif, Fernanda Viégas, Hal Abelson, James Wexler, Kristen Olson, Lucas Dixon, Mahima Pushkarna, Martin Wattenberg, Michael Terry, Rebecca Salois, Timnit Gebru, Tulsee Doshi, Yannick Assogba, Yoni Halpern, Zan Armstrong, and my other colleagues at Google for their help with this piece.

Silhouettes from ProPublica's Wee People.

More Explorables

Hidden Bias

Models trained on real-world data can encode real-world bias. Hiding information about protected classes doesn't always fix things — sometimes it can even hurt.

Modeling College GPA

Let's pretend we're college admissions officers trying to predict the GPA students will have in college (in these examples we'll use simulated data).

One simple approach: predict that students will have the same GPA in college as they did in high school.

This is at best a very rough approximation, and it misses a key feature of this data set: students usually have better grades in high school than in college

We're over-predicting college grades more often than we under-predict.

Predicting with ML

If we switched to using a machine learning model and entered these student grades, it would recognize this pattern and adjust the prediction.

The model does this without knowing anything about the real-life context of grading in high school versus college.

Giving the model more information about students increases accuracy more...

...and more.

Models can encode previous bias

All of this sensitive information about students is just a long list of numbers to model.

If a sexist college culture has historically led to lower grades for female students, the model will pick up on that correlation and predict lower grades for women.

Training on historical data bakes in historical biases. Here the sexist culture has improved, but the model learned from the past correlation and still predicts higher grades for men.

Hiding protected classes from the model might not stop discrimination

Even if we don't tell the model students' genders, it might still score female students poorly.

With detailed enough information about every student, the model can still synthesize a proxy for gender out of other variables.

Including a protected attribute may even decrease discrimination

Let's look at a simplified model, one only taking into account the recommendation of an alumni interviewer.

The interviewer is quite accurate, except that they're biased against students with a low household income.

In our toy model, students' grades don't depend on their income once they're in college. In other words, we have biased inputs and unbiased outcomes—the opposite of the previous example, where the inputs weren't biased, but the toxic culture biased the outcomes.

If we also tell the model each student's household income, it will naturally correct for the interviewer's overrating of high-income students just like it corrected for the difference between high school and college GPAs.

By carefully considering and accounting for bias, we've made the model fairer and more accurate. This isn't always easy to do, especially in circumstances like the historically toxic college culture where unbiased data is limited.

And there are fundamental fairness trade-offs that have to be made. Check out the Measuring Fairness explorable to see how those tradeoffs work.

Adam Pearce // May 2020