GeoLAN: Geometric Learning of Latent Explanatory Directions in Large Language Models

TL;DR

GeoLAN introduces a geometric training framework for large language models that enhances interpretability and fairness by promoting diverse attention and isotropy, with benefits varying across model sizes.

Contribution

The paper presents GeoLAN, a novel geometric regularization framework for LLMs that improves interpretability and fairness without sacrificing task accuracy.

Findings

GeoLAN improves geometric metrics and reduces biases in mid-sized models.

Scale-dependent trade-offs exist between geometric precision and performance.

GeoLAN maintains task accuracy while enhancing interpretability.

Abstract

Large language models (LLMs) demonstrate strong performance, but they often lack transparency. We introduce GeoLAN, a training framework that treats token representations as geometric trajectories and applies stickiness conditions inspired by recent developments related to the Kakeya Conjecture. We have developed two differentiable regularizers, Katz-Tao Convex Wolff (KT-CW) and Katz-Tao Attention (KT-Attn), that promote isotropy and encourage diverse attention. Our experiments with Gemma-3 (1B, 4B, 12B) and Llama-3-8B show that GeoLAN frequently maintains task accuracy while improving geometric metrics and reducing certain fairness biases. These benefits are most significant in mid-sized models. Our findings reveal scale-dependent trade-offs between geometric precision and performance, suggesting that geometry-aware training is a promising approach to enhance mechanistic interpretability.