Characterizing Large Language Model Geometry Solves Toxicity Detection and Generation (2312.01648v2)

Abstract: LLMs~(LLMs) drive current AI breakthroughs despite very little being known about their internal representations, e.g., how to extract a few informative features to solve various downstream tasks. To provide a practical and principled answer, we propose to characterize LLMs from a geometric perspective. We obtain in closed form (i) the intrinsic dimension in which the Multi-Head Attention embeddings are constrained to exist and (ii) the partition and per-region affine mappings of the per-layer feedforward networks. Our results are informative, do not rely on approximations, and are actionable. First, we show that, motivated by our geometric interpretation, we can bypass Llama$2$'s RLHF by controlling its embedding's intrinsic dimension through informed prompt manipulation. Second, we derive $7$ interpretable spline features that can be extracted from any (pre-trained) LLM layer, providing a rich abstract representation of their inputs. Those features alone ($224$ for Mistral-7B/Llama$2$-7B and $560$ for Llama$2$-70B) are sufficient to help solve toxicity detection, infer the domain of the prompt, and even tackle the Jigsaw challenge, which aims at characterizing the type of toxicity of various prompts. Our results demonstrate how, even in large-scale regimes, exact theoretical results can answer practical questions in LLMs. Code: \url{https://github.com/RandallBalestriero/SplineLLM}.

Characterizing LLMs Geometry Solves Toxicity Detection and Generation

The paper "Characterizing LLMs Geometry Solves Toxicity Detection and Generation" provides a geometric analysis of LLMs, specifically focusing on how these models represent data internally and how such understanding can be leveraged to improve practical applications like toxicity detection. The authors present a detailed paper on the geometric structure of LLMs’ computations, introducing new insights and practical solutions for downstream tasks which are increasingly relevant in human-computer interaction scenarios.

Geometric Insights into LLMs

At the core of this research lies the characterization of LLMs through a geometric perspective. Particularly, the paper explores the workings of the Multi-Head Attention (MHA) mechanism and feedforward networks in LLM layers. The authors provide exact formulations for the intrinsic dimension of MHA embeddings, suggesting these embeddings reside on manifolds defined by the token sequences’ convex hulls. The intrinsic dimension is critical as it dictates the expressiveness and capacity of the LLMs. This intrinsic dimension is manipulated through prompt engineering to control downstream generation tasks, including bypassing reinforcement learning from human feedback (RLHF).

Intrinsic Dimension and Toxicity Generation

The work demonstrates how intrinsic dimensionality, determined by input token interdependencies captured in the attention matrices, impacts the model's output expressivity. A key finding is the ability to increase this dimension through informed prompt manipulation, thereby exploring latent spaces within the model not adjusted by RLHF. The authors illustrate that feeding LLMs with related sentences boosts the intrinsic dimensionality, sometimes resulting in a breach of RLHF protections leading to toxic content generation, highlighting potential vulnerabilities in current RLHF practices.

Spline Operators and Feature Engineering

The authors further explore MLP layers in LLMs by leveraging spline theory to express feedforward computations as Continuous Piecewise Affine (CPA) mappings. This perspective provides a foundational understanding of how LLMs partition the input space and allows for the derivation of seven spline features that effectively capture the geometrical properties of prompt representations. These features encapsulate the model’s response to input across layers, proving robust in discriminating between different input domains and effectively detecting text toxicity.

Practical Implications and Evaluation

Through extensive empirical validation, the authors demonstrate the utility of the geometric features in tasks like toxicity detection. The derived features allow for competitive performance in separating toxic from non-toxic prompts, markedly outperforming commercially available solutions in efficiency and accuracy. Notably, these features when used with simple linear models, achieve high ROC-AUC scores on standard datasets such as Jigsaw’s toxic comment classification challenge, showcasing near-exhaustive coverage of toxic modalities without additional data augmentation or model fine-tuning.

Conclusions and Future Directions

This paper contributes significantly to the field of AI interpretability and application, providing both theoretical insights and practical tools for improvement in toxicity detection and model understanding. The ability to manipulate the expressivity of LLMs through intrinsic geometric features opens avenues for safer AI applications and highlights potential areas for future research. Ensuring robust and comprehensive safety mechanisms against toxicity requires further explorations into geometric interpretations of model structures and beyond RLHF engagements.

Future work may involve extending these geometric analyses to larger and more complex models, potentially uncovering novel features that further enhance model transparency and control. Additionally, integrating this understanding with diverse datasets can help generalize these findings, improving the adaptability and robustness of models in varied real-world scenarios.