Growing Prototypes on Demand: Scalable Prototype Learning for Causal Language Models

Enabling causal language models to produce interpretable evidence remains challenging. Prototype-based reasoning offers “this looks/reads like that" interpretability, but compressing diverse contextual training evidence into reusable prototypes remains challenging at scale. We propose a non-parametric approach that frames prototype learning as clustering in a frozen model's hidden space using a Dirichlet Process Gaussian Mixture Model (DP-GMM), which allows the number of components to grow only when needed. To scale to millions of hidden states, we introduce a tree-structured DP-GMM that performs coarse-to-fine mixture fitting, with a separate prototype set and prediction head fitted for each task and training distribution on a frozen backbone. The expectations of leaf means are projected to training hidden states to obtain example-grounded prototypes, whose normalised distance similarities drive a learned prediction head. This design enables scalable DP-GMM fitting with node-local cost independent of global prototype capacity for fixed local sample size, hidden dimension, truncation, and fitting settings. Experiments across language modelling, captioning, and VQA fit approximately 2.3M-7.4M hidden states into 5.8k-7.6k prototypes, with small-to-moderate performance drops. Paired interventions show that retaining a small top-matched set largely preserves performance, while removing it severely degrades performance, which supports the computational faithfulness.