Eliciting Latent Predictions from Transformers with the Tuned Lens

“7B prompted with an except from the abstract of Vaswani et al. (2017).”

“The impressive performance of transformers in natural language processing (Brown et al., 2020) and computer vision (Dosovitskiy et al.”

“The impressive performance of transformers in natural language processing (Brown et al., 2020) and computer vision (Dosovitskiy et al., 2020) suggests that their internal representations have rich structure worthy of scientific investigation.”

“…flatten these trajectories into feature vectors and feed them into two standard outlier detection algorithms: isolation forest (iForest) (Liu et al., 2008) and local outlier factor (LOF) (Breunig et al., 2000), both implemented in scikit-learn (Pedregosa et al., 2011) with default hyperparameters.”

“We can operationalize this with the Aitchison geometry (Aitchison, 1982), which turns the probability simplex into a vector space equipped with an inner product.”

“6 We then flatten these trajectories into feature vectors and feed them into two standard outlier detection algorithms: isolation forest (iForest) (Liu et al., 2008) and local outlier factor (LOF) (Breunig et al., 2000), both implemented in scikit-learn (Pedregosa et al., 2011) with default…”

“This can be viewed as a distillation loss, using the final layer distribution as a soft label (Sanh et al., 2019).”

“, 2008) and local outlier factor (LOF) (Breunig et al., 2000), both implemented in scikit-learn (Pedregosa et al.”

“Stochastic depth (Huang et al., 2016), also known as LayerDrop (Fan et al., 2019), is a regularization technique which randomly drops layers during training, and can be applied to both ResNets and transformers.”

“One simple technique is to fit a multivariate Gaussian to the model’s final layer hidden states on the training set, and flag inputs as OOD if a new hidden state is unusually far from the training distribution as measured by the Mahalanobis distance (Lee et al., 2018; Mahalanobis, 1936).”

“The beliefs of a rational agent should not update in an easily predictable direction over time (Yudkowsky, 2007), since predictable updates can be exploited via Dutch books (Ramsey, 1926; De Finetti et al., 1937; Hacking, 1967; Garrabrant et al., 2016).”

“Even when the logit lens is useful, we find that it is a biased estimator of the model’s final output: it systematically puts more probability mass on certain vocabulary items than the final layer does.”

“They then sampled from the vectors that scored < 0.1 toxicity (where a score > 0.5 is classified as toxic), set each value vector’s respective coefficient to 3 (because 3 was a higher than average activation), and ran the newly edited model through a subset of REALTOXICPROMPTS (Gehman et al., 2020) to measure the chan…”

“One common approach is to train classifiers to extract specific concepts from hidden states, like part-of-speech and syntactic structure (Hewitt and Manning, 2019; Tucker et al., 2021; Li et al., 2022).”

“Composing a translator with the pretrained unembedding yields a probe (Alain and Bengio, 2016) that maps a hidden state to a distribution over the vocabulary.”

“Since BLOOM and GPT-2 do not have publicly available validation sets, we use the Pile validation set (Gao et al., 2020; Biderman et al., 2022).”

“The performance of both lenses was then evaluated on the test set of Anthropic’s Helpful Harmless (Bai et al., 2022) conversation dataset and the RedPajama Dataset.”

“The iteration learned is defined as the earliest training step τ where the model’s top-1 prediction for a datapoint x is fixed (Toneva et al., 2018).”

“This is especially important for probing techniques, since it’s known that probes can learn to rely on spurious features unrelated to the model’s performance (Hewitt and Liang, 2019; Be-linkov, 2022).”

“Stochastic depth (Huang et al., 2016), also known as LayerDrop (Fan et al., 2019), is a regularization technique which randomly drops layers during training, and can be applied to both ResNets and transformers.”

“An affine transform suffices to stitch together independently trained models with minimal performance loss (Bansal et al., 2021; Csiszárik et al., 2021).”

“…that “write" (in the sense of Elhage et al. (2021)) to intermediate hidden states of a language model (Millidge and Black, 2022; Dar et al., 2022; Geva et al., 2022; 2020), and even use this analysis to successfully edit model behavior (Geva et al., 2022; Millidge and Black, 2022; Dai et al.,…”

“, 2022) and DeeBERT (Xin et al., 2020) are based on the observation that “easy” examples require less computation to classify than “difficult” examples.”

“This is especially important for probing techniques, since it’s known that probes can learn to rely on spurious features unrelated to the model’s performance (Hewitt and Liang, 2019; Belinkov, 2022).”

“Prior work has argued that interpretability hypotheses should be tested with causal experiments: an interpretation of a neural network should make predictions about what will happen when we intervene on its weights or activations (Olah et al., 2020; Chan et al., 2022).”

“Following Baldock et al. (2021)’s work on computer vision models, we define the prediction depth of a prompt x to be the number of layers after which a model’s top-1 prediction for x stops changing.”

“In particular, we use the pretrained BPEmb encoder (Heinzerling and Strube, 2018) 9 .”

“One common approach is to train classifiers to extract specific concepts from hidden states, like part-of-speech and syntactic structure (Hewitt and Manning, 2019; Tucker et al., 2021; Li et al., 2022).”

“These results strongly support the idea that adjacent layers in a ResNet encode fundamentally similar representations (Greff et al., 2016).”

“…that “write" (in the sense of Elhage et al. (2021)) to intermediate hidden states of a language model (Millidge and Black, 2022; Dar et al., 2022; Geva et al., 2022; 2020), and even use this analysis to successfully edit model behavior (Geva et al., 2022; Millidge and Black, 2022; Dai et al.,…”

“We find that the tuned lens resolves the problems with the logit lens discussed in Section 2: it has significantly lower 3 Pythia and GPT-NeoX-20B were trained using the same architecture, data, and codebase (Andonian et al., 2021).”

“Amnesic probing (Elazar et al., 2021) provides one way to do this—it seeks a direction whose erasure maximally degrades a model’s accuracy.”

“But recent work has pointed out that setting activations to zero may inadvertently push a network out of distribution, making the experimental results hard to interpret (Wang et al., 2022; Chan et al., 2022).”

“Then, the mutual similarity of words can easily be measured by their cosine similarity under a word embedding (Sitikhu et al., 2019).”

“1), we find that tuned lens prediction trajectories can be used to detect prompt injection attacks (Perez and Ribeiro, 2022) often with near-perfect accuracy (Section 5.”

“An affine transform suffices to stitch together independently trained models with minimal performance loss (Bansal et al., 2021; Csiszárik et al., 2021).”

“To address the shortcomings of the logit lens, we introduce the tuned lens .”

“The idea of model stitching was introduced by Lenc and Vedaldi (2015), who form a composite model out of two frozen pretrained models A and B , by connecting the bottom layers of A to the top layers of B .”

“Jastrz˛ebski et al. (2017) argue that skip connections encourage neural networks to perform “iterative inference” in the following sense: each layer reliably updates the hidden state in a direction of decreasing loss.”

“Another issue with the logit lens is that transformer hidden states often contain a small number of very high variance dimensions, and these “rogue dimensions” (Timkey and van Schijndel, 2021) tend to be distributed unevenly across layers; see Figure 6 (top) for an example.”

“We leave out the final transformer layer for GPT-2 (Radford et al., 2019), GPT-NeoX-20B (Black et al., 2022), OPT (Zhang et al., 2022), and Pythia (Biderman et al., 2023), and include it for GPT-Neo (Black et al., 2021).”

“One common approach is to train classifiers to extract specific concepts from hidden states, like part-of-speech and syntactic structure (Hewitt and Manning, 2019; Tucker et al., 2021; Li et al., 2022).”

“However, this only elicits a single important direction, (4)For simplicity we assume the “features as directions” hypothesis (Elhage et al., 2022), which defines a “feature” to be the one-dimensional subspace spanned by a unit vector v.”

“Ablating an outlier direction can drastically harm performance (Kovaleva et al., 2021), so if LogitLens relies on the presence or absence of particular outlier dimensions, the perplexity of logit lens predictions might be spuriously high.”

“In order to downweight the influence of rare tokens, we use the weighted Aitchison inner product introduced by Egozcue and Pawlowsky-Glahn (2016), defined as where w is a vector of positive weights, and g w ( p ) is the weighted geometric mean of the entries of p .”

“We start by extending a recent use case of the logit lens.”

“An affine transform suffices to stitch together independently trained models with minimal performance loss (Bansal et al., 2021; Csiszárik et al., 2021).”

“Even when the logit lens is useful, we find that it is a biased estimator of the model’s final output: it systematically puts more probability mass on certain vocabulary items than the final layer does.”

“One common approach is to train classifiers to extract specific concepts from hidden states, like part-of-speech and syntactic structure (Hewitt and Manning, 2019; Tucker et al., 2021; Li et al., 2022).”

“Since BLOOM and GPT-2 do not have publicly available validation sets, we use the Pile validation set (Gao et al., 2020; Biderman et al., 2022).”

“Halawi et al. (2023) apply the logit lens to downstream tasks with few-shot prompts, and find that predictions extracted from earlier layers are often more robust to incorrect demonstrations than final layer predictions.”

“While they’re not officially the same model suite, they’re more consistent than the OPT models. bias (Figure 3), and much lower perplexity than the logit lens across the board (Figure 5, Appendix A).”

“We leave out the final transformer layer for GPT-2 (Radford et al., 2019), GPT-NeoX-20B (Black et al., 2022), OPT (Zhang et al., 2022), and Pythia (Biderman et al., 2023), and include it for GPT-Neo (Black et al., 2021).”

“We trace the flow of numerical information within Transformer-based LMs: given an input query, the model processes the representations of numbers and operators with early layers (A).”

“Unlike fixed injection templates used in [30] (e.g., “ignore previous prompts”), our method adapts to diverse scenarios of injection forms.”

“We also aim to test newer probing techniques like the tuned lens (Belrose et al., 2023) in future research.”

“Other related efforts have passed hidden intermediate states directly to the decoder head (skipping in-between layers) to “verbalize” such embeddings (Din et al., 2023; Belrose et al., 2023; nostalge-braist, 2020).”

“A critical tool in these interpretability workflows is the ability to hook the output of intermediate layers to understand what the model may be “thinking” as it refines the final output [3, 40, 51].”

“This hypothesis could, for instance, be tested by applying the tuned lens method (Belrose et al., 2023), which reveals model predictions about upcoming input at different LLM layers, other than just the output layer.”

“…computational exper-iment, we just analyzed the probability output of a target word through the Surprisal scores, but in the future, it would be useful to adopt interpretability techniques to get more insights on the hidden representations of the NLMs (Yin and Neubig, 2022; Belrose et al., 2023).”

“Adapting interpretability techniques from Tuned Lens, 22 Captum, 23 DeepLift, 24 and SHAP, 25 we systematically decoded how specific amino acid residues, their patterns, and positional relationships influence model decisions (Figures 3, 4, 5, 6, 7, and S4–S6).”

“Similarly, the tuned lens (Belrose et al., 2023) learns affine transformations that translate representations from an internal layer to the final layer.”

“While not exclusively mechanistic, several other recent researches have explored the hidden representations and the functioning of the inner components of large Language Models [20, 21].”

“Targeted interventions on the computation are then applied to estimate the impact of individual components on model behavior (Modarressi et al., 2022; Mohebbi et al., 2023; Wang et al., 2022; Nanda et al., 2023; Merullo et al., 2023; Belrose et al., 2023; Meng et al., 2023a; Vig et al., 2020).”

“…have shown interesting insights into how LLMs learn new concepts throughout their training and how different layers learn differently [2], which can be used for proposing meta frameworks such as ours by observing which layers of a model might focus on which domain of knowledge…”

“…for GenAIware (Section 3.1), then, we focus on the existing frameworks designed to manage the interaction with LLMs (Section 3.2), and on the known challenges arising from GenAIware development, and, finally, we position the proposed LSL framework with respect to the state-of-the-art (Section 3.4).”