Characterizing Large Language Model Geometry Helps Solve Toxicity Detection and Generation

Abstract: LLMs drive current AI breakthroughs despite very little being known about their internal representations. In this work, we propose to shed the light on LLMs inner mechanisms through the lens of geometry. In particular, we develop 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 feedforward (MLP) network of LLMs' layers. Our theoretical findings further enable the design of novel principled solutions applicable to state-of-the-art LLMs. First, we show that, through our geometric understanding, we can bypass LLMs' RLHF protection by controlling the embedding's intrinsic dimension through informed prompt manipulation. Second, we derive interpretable geometrical features that can be extracted from any (pre-trained) LLM, providing a rich abstract representation of their inputs. We observe that these features are sufficient to help solve toxicity detection, and even allow the identification of various types of toxicity. Our results demonstrate how, even in large-scale regimes, exact theoretical results can answer practical questions in LLMs. Code: https://github.com/RandallBalestriero/SplineLLM

• The paper reveals that LLMs possess intrinsic dimensions within attention embeddings, enabling refined control over toxicity generation.
• It employs spline theory to model MLP layers as continuous piecewise affine mappings for effective feature extraction and toxicity detection.
• Empirical evaluations show these geometric features outperform mainstream RLHF models in accuracy and robustness on toxicity classification benchmarks.

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 study 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.