Prototype Language Models

Jul 16, 2026

Upload by Himbo

Knowing which training examples drive model outputs is crucial for auditing, correcting, and understanding modern LLM behavior, yet existing language-model explanations remain expensive, approximate, and largely post-hoc. This work proposes PRISM, a prototype language-model architecture that predicts via a sparse, non-negative mixture of learned prototypes. Clustering objectives anchor prototypes to coherent neighborhoods in training data. Across 130M–1.6B parameters and up to 50B tokens, PRISM matches or exceeds dense baselines, localizes curvature for tractable Hessians, and enables fast attribution (~500×).

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages54

Reading Time2 h 16 min

Introduction

PRISM: Prototypes for Interpretable Sequence Modeling

Warm-up: reviving ProtoPNet for language modeling

Sparse prototypical reconstruction

Automated interpretability pipeline

Using TinyStories as a Microscope

Anatomy of a prototype prediction

Training and interpretability dynamics

Prototype structure is learned into the backbone

Training Data Attribution in a Prototype Subspace

Clustering localizes curvature in prototype space

Cacheable prototype space influence functions

Cached attribution preserves signal at lower cost

Scaling Prototype Language Models

Scaling setup and downstream performance

Training stability and efficiency

New Workflows Enabled by Prototypes

Understanding model behavior through prototypes

Prototype controllers boost performance and trace corrections to training data

Preference alignment without finetuning

Limitations and Future Roadmap

Sequence and document attribution

Hessian-aware model design

Deep prototype language models

Retraining and learned attribution objectives

Related work

Conclusion

arXiv :2607 .005 10v 1 [ cs .LG] 1 Jul 2026 Prototype Language Models Dan Ley1,2* Giang Nguyen2 Himabindu Lakkaraju1 Julius Adebayo2 1 Harvard University 2 Guide Labs Inc. * Work initiated during an internship at Guide Labs. Abstract Knowing which training examples drive outputs is fundamental to auditing, correcting, and understanding language models, yet for modern LLMs this remains expensive, approximate, and largely post-hoc. Standard language models generate tokens through a dense network pathway, causing training data's influence to be distributed across parameters rather than organized along explicit, traceable components. We introduce a prototype language model architecture, Prototypes for Interpretable Sequence Modeling (PRISM), that forms each prediction via a sparse, non-negative mixture of learned prototypes, trained with clustering objectives that anchor each prototype to coherent neighborhoods of training examples. Across architectures from 130M to 1.6B parameters trained on up to 50B tokens, prototype language models either surpass or remain within 2.5 percentage points on average downstream accuracy of matched dense baselines. We show that sparse prototype structure localizes curvature in the loss landscape, yielding a more tractable Hessian and enabling training data attribution that is ∼500 × faster than post hoc baselines when consuming equivalent memory. Calibrating linear prototype controllers can improve downstream accuracy by roughly 3 points while tracing those corrections back to training neighborhoods, and targeted prototype suppression can remove model behaviors without finetuning or measurable loss in generation quality. Standard language models entangle training effects, while PRISM makes them traceable through prototypes. Contents 1 Introduction 1 2 PRISM: Prototypes for Interpretable Sequence Modeling 3 2.1 Warm-up: reviving ProtoPNet for language modeling ..................... 3 2.2 Sparse prototypical reconstruction ................................. 4 2.3 Automated interpretability pipeline ................................. 6 3 Using TinyStories as a Microscope 6 3.1 Anatomy of a prototype prediction ................................. 7 3.2 Training and interpretability dynamics ............................... 8 3.3 Prototype structure is learned into the backbone ........................ 10 4 Training Data Attribution in a Prototype Subspace 11 4.1 Clustering localizes curvature in prototype space ........................ 12 4.2 Cacheable prototype space influence functions ......................... 14 4.3 Cached attribution preserves signal at lower cost ........................ 15 5 Scaling Prototype Language Models 17 5.1 Scaling setup and downstream performance ........................... 17 5.2 Training stability and efficiency .................................... 19 6 New Workflows Enabled by Prototypes 19 6.1 Understanding model behavior through prototypes ...................... 20 6.2 Prototype controllers boost performance and trace corrections to training data .... 21 6.3 Preference alignment without finetuning ............................. 23 7 Limitations and Future Roadmap 24 7.1 Sequence and document attribution ................................ 25 7.2 Hessian-aware model design ..................................... 25 7.3 Deep prototype language models .................................. 25 7.4 Retraining and learned attribution objectives ........................... 26 8 Related work 26 9 Conclusion 28 1 Introduction When a language model (LM) produces a harmful response, copyrighted content, or makes an inaccurate factual claim, a natural follow-up question is: which training data made that output likely? Training data attribution (TDA) consists of a family of techniques to answer this question, characterizing how individual examples, groups of examples, or broader training sources influence model predic

FAQs

What problem does the paper target in language model interpretability?

It targets training data attribution: identifying which training examples most likely caused a specific model output, despite standard LLMs entangling effects across many parameters and making direct tracing difficult.

How does PRISM make predictions and improve interpretability?

PRISM forms each next-token prediction via a sparse, non-negative mixture of learned prototypes. Prototype clustering objectives anchor each prototype to coherent neighborhoods of training examples, making effects traceable through the prototype structure.

What evidence is given that prototype models are effective and efficient?

Across model sizes (130M to 1.6B) and up to 50B tokens, prototype language models either match or improve downstream accuracy versus dense baselines by about 2.5 percentage points. The sparse structure localizes curvature, yields a more tractable Hessian, and enables training-data attribution roughly 500× faster than post-hoc baselines at comparable memory use.

Jul 16, 2026

Upload by Himbo