PaperlayerView exact revision on arXiv

Source-linked AI summary

Inspecting and Editing Knowledge Representations in Language Models

Evan Hernandez, Belinda Z. Li, Jacob Andreas

arXiv:2304.00740v3cs.CL

TL;DR

Language models encode factual knowledge from training and input context, but these representations can produce incorrect or incoherent generations. REMEDI learns textual-to-representation mappings that edit or probe those encodings, improving factual control while also revealing when outputs may conflict with represented knowledge.

  • Problem

    Language models can contain mistaken factual encodings and fail to integrate contextual information, creating a need to inspect and control the representations underlying generation.

  • Method

    REMEDI learns local transformations from textual attribute descriptions and applies them to entity representations, while comparing edit vectors with representations to probe model knowledge.

  • Results

    REMEDI improves factual editing and generalizes across background and in-context knowledge, increasing GPT-J accuracy by over 15% on held-out entity–attribute pairs while preserving fluency.

  • Takeaways & Limitations

    Directly specifying facts in an LM’s native representation scheme can support controlled generation, factual-error detection, and inspection of model knowledge.

  • Takeaways & Limitations

    REMEDI requires fact encodings that are causally linked to generation, since information present in representations may not influence subsequent outputs.

Abstract

from arXiv · show

Neural language models (LMs) represent facts about the world described by text. Sometimes these facts derive from training data (in most LMs, a representation of the word "banana" encodes the fact that bananas are fruits). Sometimes facts derive from input text itself (a representation of the sentence "I poured out the bottle" encodes the fact that the bottle became empty). We describe REMEDI, a method for learning to map statements in natural language to fact encodings in an LM's internal representation system. REMEDI encodings can be used as knowledge editors: when added to LM hidden representations, they modify downstream generation to be consistent with new facts. REMEDI encodings may also be used as probes: when compared to LM representations, they reveal which properties LMs already attribute to mentioned entities, in some cases making it possible to predict when LMs will generate outputs that conflict with background knowledge or input text. REMEDI thus links work on probing, prompting, and LM editing, and offers steps toward general tools for fine-grained inspection and control of knowledge in LMs.

1 Introduction

REMEDI addresses errors in language-model representations by learning to inspect and modify factual encodings directly. It supports both controlled generation and prediction of potentially incorrect or incoherent outputs.

  • Motivation: Language models encode general knowledge and situations described in input text, but erroneous, outdated, contradictory, or biased training information can produce generation errors.Understanding how internal representations relate to output also matters for attribution and controlled generation.
  • Illustration: Figure 1 shows that REMEDI can patch LM errors, insert facts with or without prompt context, and help detect errors before generation.
  • REMEDI: REMEDI learns directions in internal representation spaces corresponding to factual attributes and adds them to entity representations to edit attributed facts.These edits can sometimes produce outputs unavailable through a corresponding textual prompt.
  • REMEDI: REMEDI encodings can probe factual knowledge and predict when models will generate incorrect or incoherent output.
  • Implications: REMEDI generalizes from modifying background knowledge to querying and modifying knowledge specified in context.The findings suggest that models represent and integrate these information sources in a unified manner.

2 REMEDI

REMEDI learns local interventions that transform an entity’s hidden representation so generated text reflects a desired attribute. It also supports out-of-context control and representation-based inspection while preserving intermediate behavior through non-destructive training objectives.

  • Motivation: REMEDI targets context-integration failures and unwanted prior associations by changing an entity’s internal encoding rather than relying solely on textual prompts.Prompt engineering can be opaque, non-systematic, and labor-intensive when a suitable prompt is difficult to find.
  • Overview: The method applies a single-vector update to one entity token at one layer, enabling controlled generation without textual prompts and comparison-based inspection of entity representations.
  • Editing representations: REMEDI assumes autoregressive transformer language models and modifies hidden states to change the probability of subsequent tokens.
  • Editing representations: Entity representations are suitable intervention sites because factual and semantic properties have been shown to be linearly decodable from them.The method therefore transforms the representation of the entity token to encode a target fact.
  • Design requirement: An effective editor must target fact encodings that causally influence generation, because decodable information may not affect subsequent outputs.
  • Illustration: Figure 2 illustrates editing John’s representation with the attribute “plays the oboe” so an appropriate completion becomes “the concert hall.”
  • REMEDI EDITING: The editor is trained on contexts, entity indices, attributes, and target completions to maximize the target completion probability after intervention.The target attribute representation is obtained by averaging the model’s encoding of the attribute text.
  • REMEDI EDITING: Non-destructive editing suppresses a strong prior association while penalizing changes to intermediate-token distributions, preserving behavior before the target prediction.The complete objective combines these terms with hyper-parameters λ1 and λ2.

3 Related Work

Related work studies how language models encode factual and context-dependent knowledge, adapt through prompts or learned tokens, and localize or edit facts in model parameters. REMEDI instead intervenes directly in contextual representations.

  • Probing factual knowledge: Prior probing work extracts context-agnostic factual knowledge from language models through textual queries designed for different knowledge sources.
  • Context-dependent representations: Other work shows that language models represent context-dependent entity states over discourse and evaluates this capability through downstream reading-comprehension tasks.
  • Editing language models: Language-model adaptation methods include textual prompt design and prefix-tuning with learned continuous tokens prepended to examples.
  • Editing language models: Parameter-editing methods localize and modify factual information in model weights, whereas REMEDI operates on representations rather than weight matrices.

4 Controlling Generation

REMEDI controls generation by inserting learned attribute encodings into entity representations, both to patch context-integration failures and to overwrite background facts. Across occupation classification and COUNTERFACT editing, it improves factual consistency while generally preserving fluency, entity identity, and locality.

  • Patching Errors: GPT-J achieves 55% accuracy on the in-context occupation task, while REMEDI increases held-out accuracy by over 15% and preserves fluency.The task ranks 28 occupations, with correctness determined by whether the true occupation is ranked first.
  • Method: REMEDI learns entity representations encoding target attributes and applies each edit to one token at one hidden layer.The editor can be trained out-of-context and evaluated with or without additional textual context.
  • Patching Errors: REMEDI produces fluent text consistent with inserted occupations even without textual context, slightly outperforming its in-context editing effectiveness despite out-of-context training.This indicates that representation edits can replace textual prompting for supplying occupation information.
  • Editing Factual Associations: Unlike global parameter-editing methods, REMEDI limits changes to the entity of interest and avoids neighboring-entity failures that occur about 21% of the time with ROME.REMEDI never causes the neighboring-entity failures reported for ROME in this comparison.
  • Editing Factual Associations: On COUNTERFACT, REMEDI is comparably effective to model-editing methods and more effective than prefixing prompts with the new fact.The benchmark evaluates efficacy, neighborhood, consistency, fluency, and essence across edited factual associations.
  • Editing Factual Associations: REMEDI incorporates new information 98.2% of the time versus 80.2% for the base LM, though this comes at some cost to entity essence.The Euler example illustrates that edits may alter correlated properties while respecting implicit associations.

5 Detecting Errors

REMEDI probes LM representations to detect factual errors in background knowledge and failures to integrate contextual information. It outperforms comparison methods in both settings, although contextual-error detection remains less accurate than direct supervision.

  • Method: The probe detects whether an LM has acquired background or contextual knowledge by inspecting representations for information REMEDI would add.This extends editing encodings into an evaluation tool for unmodified models.
  • Method: REMEDI compares entity representations with encodings of true and distractor attributes to predict when an LM will generate an incorrect output.For example, it computes encodings for “located in London” and distractor attributes such as “located in Arizona.”
  • Detecting errors in prior knowledge: REMEDI outperforms all methods except the task-trained shortcut model on factual error detection using both F1 and ϕ.The evaluation includes baselines, controls, a random model, and a supervised error skyline.
  • Predicting errors in context: REMEDI outperforms all baselines, including shortcut, when detecting contextual errors.The method was trained for editing in non-contextual sentences, yet its encodings align with knowledge provided in context.
  • Limitations: Contextual-error detection is not extremely accurate and is substantially worse than a model directly supervised on LM error information.This supervised skyline is not directly comparable because it is trained on annotated model outputs.

6 Conclusions

The paper presents REMEDI as a way to interpret and control factual knowledge in LMs through local transformations of entity representations.

  • Conclusion: REMEDI constructs local transformations from textual attribute descriptions and applies them to contextual representations of entity mentions and other nouns.Amplifying a fact’s encoding can force generation consistent with that fact, even when a textual prompt does not.

Ethical Considerations

Controlling LM generations can reduce harmful or misleading outputs but can also help malicious actors produce them.

  • Ethical considerations: Stronger control of language generation is both beneficial for preventing harmful or misleading outputs and risky because it can encourage such outputs.The paper argues that uncontrolled factual or incoherent generations pose risks when they reach users through trusted applications.

Reproducibility Statement

The paper states that its code, data, library, figure-generation code, and experiment details will be made publicly available.

  • Reproducibility: All code and data, including the REMEDI library and figure-generation code, will be publicly available upon publication.The paper also describes experiment details, preprocessing procedures, and hyperparameter sweeps.

A Limitations

REMEDI’s production use is constrained by its reliance on in-domain training data and simplified prompting settings. Detecting certain failures also requires prior knowledge of the correct attribute and a distractor for comparison.

  • REMEDI’s linear editing functions require in-domain training data formatted with a prompt, entity, attribute, and target word.This requirement may restrict deployment when suitable training data are unavailable.
  • Failure detection requires knowing the correct attribute beforehand and having a distractor attribute for comparison.Neither resource may be available in practice.
  • The prompting settings are deeply simplified for controlled experimentation, limiting how directly the evaluations represent production use.The paper notes that the benchmark prompts often work without REMEDI in state-of-the-art models, while failures arise in more complex prompts.

B Dataset Preprocessing

The experiments preprocess each dataset into structured context, entity, attribute, and target examples, with separate training and held-out samples. McRae Norms additionally supplies human feature associations and probability-based filtering.

  • COUNTERFACT: COUNTERFACT uses the first paraphrase prompt, retains the entity sentence, and treats every token after the entity as the attribute.Objectives and evaluations use the record’s primary prompt.
  • Bias in Bios: Bias in Bios uses the second bio sentence longer than three words as context and normalizes entity mentions before defining the occupation prompt and target.The processed data are randomly split into 5000 training and 5000 held-out records.
  • McRae Norms: McRae Norms represents each concept with human-associated features and probabilities reflecting how many of thirty people attributed each feature.Feature pairs are sampled when they co-occur for at least one concept.
  • McRae Norms: The preprocessing estimates feature co-occurrence using p(f2 | c)p(f1 | c) across concepts under uniform p(c) and conditional independence given c.The resulting human-derived probabilities support evaluation and candidate-pair filtering.
  • Editor data: REMEDI training uses 5000 sampled pairs, while another 5000 are held out; concepts lacking both features provide context and test prompts.The editor maximizes the probability of a heuristically selected final prompt token.

C Training Editors

The editors are trained with dataset-specific objectives and selected by layer sweeps, while concept-editing evaluations test whether REMEDI changes related features selectively. Results indicate stronger effects than prefixing and structured changes among original features.

  • Training: Editors are trained on 5000 examples with 500 validation samples using AdamW, dataset-specific loss weights, and early stopping after two unimproved validation epochs.COUNTERFACT uses λ1 = 1 and λ2 = 10, whereas Bias in Bios and McRae Norms omit the prior term.
  • Layer selection: Layer sweeps evaluate editors across GPT-J layers using generation metrics on 1000 held-out records before selecting the editing layer.Greedy decoding is used for sweeps and top-k sampling with k = 5 for final evaluations.
  • Layer selection: COUNTERFACT performs best at earlier layers, whereas Bias in Bios and McRae Norms are relatively flat across early and middle layers.The selected layers are 1 for COUNTERFACT and 12 for Bias in Bios and McRae.
  • Concept editing: Compared with prefixing, REMEDI produces an order of magnitude larger probability increase for correlated features and stronger alignment with human feature co-occurrences.The comparison uses ΔpLM and r(pLM, pH).
  • Concept editing: REMEDI changes original features in both directions, decreases their human-distribution correlation, and has negligible effects on unrelated features.These effects are consistent with promoting compatible features and suppressing features that conflict with the added attribute.
  • Concept editing: Figure 4 illustrates that adding a chopping-wood feature raises probabilities for correlated features while original and unrelated features remain lower.The result supports global modification of concept relations rather than simple priming of a target string.

E Results in Other Models

Across GPT2-XL and Llama-2-13b, REMEDI generally preserves GPT-J’s trends, with effectiveness increasing with model size and stronger control than prompting. Results also expose model-dependent context effects and several failure modes.

  • Cross-model trends: Across GPT2-XL and Llama-2-13b, REMEDI retains GPT-J’s trends, becomes more effective with larger models, and transfers from non-contextual training to contextual classification.For generation, REMEDI provides more effective control than standard prompting.
  • Model differences: GPT2-XL follows Bias in Bios prompts poorly, while REMEDI substantially improves steering without textual context and improves the contextual baseline only slightly.Llama-2-13b follows prompts well at baseline, so in-context REMEDI yields a smaller accuracy improvement.
  • Failure modes: REMEDI can produce disfluent or incorrect generations, partial edits, damaged entity essence, or changes to unrelated facts.These failures occur in both generic and famous-entity settings.
  • Failure modes: Some errors also appear without REMEDI, and GPT-J may already encode incorrect facts; larger datasets or edits at different layers could mitigate them.The paper therefore attributes some failures partly to the underlying model rather than solely to the editor.
  • Generalization: REMEDI performs best on most metrics for attributes unseen during training, although efficacy is slightly lower and the authors suggest possible linear-editor overfitting.For unseen attributes, generations are more fluent and more essence-preserving.
  • Dependence on prior knowledge: REMEDI performs slightly better when the language model already ranks the correct fact above a distractor before editing.This analysis concerns held-out COUNTERFACT entity-attribute pairs.
  • Representation norms: REMEDI directions can greatly exceed the norm of the edited representation, sometimes making it more than twice its pre-edit size.The paper notes that representation norms differ substantially across input types, complicating interpretation of abnormality.
Loading 2304.00740v3…