Small language models cut output layer with geometry based decoding
RiLM: Parameter-Efficient Language Modeling via Geodesic Decoding
Computation and LanguageArtificial Intelligence
Summary
Small language models often waste a lot of their limited capacity on the part that predicts the next word. The authors propose a new method, called RiLM, that replaces this prediction layer with a geometric approach where the model's current state moves along a curved space, and next-word chances are based on distances in that space. They tested two versions: one flat and one curved like a ball, and the curved version performed better on predicting text. This approach is efficient and works well even with small models and limited vocabulary sizes.
What this means in practice
- •For mobile app developers: Deploy compact language models on devices with limited memory and compute by removing large output layers and using efficient geometric decoding.
- •For natural language processing engineers: Adapt language models to new domains with smaller parameter budgets while maintaining prediction quality through geometry-based decoding.
Authors
Fang Li
Abstract
Language models under one million parameters matter for edge deployment, domain adaptation, and reproducible research, yet a two-layer LSTM or Transformer at embedding width d = 128 still spends roughly one third of its capacity on the output matrix W_out in R^(d x |V|). We propose Riemannian Language Models (RiLM), which remove that layer entirely: context unfolds as a trajectory on a Riemannian manifold, and next-token probabilities arise from squared geodesic distance between the current state and vocabulary embeddings. The same embedding map serves input and output -- decoding is geometry. We instantiate the framework on flat R^d (Flat RiLM) and the Poincare ball H^d (HypRiLM) with a shared MLP composition map phi (~290k parameters, d = 128, |V| = 2000). Across five seeds on WikiText-2, HypRiLM reaches 54.2 +/- 0.2 validation perplexity versus 87.6 +/- 0.6 for Flat RiLM; tied and matched LSTM, Transformer, and SSM controls remain at 113-147 PPL on WT-2 -- HypRiLM leads by roughly 2x over the strongest tied recurrent baseline (SSM, 113.0 +/- 3.8). Penn Treebank and a 10k-vocabulary stress test confirm that geodesic decoding transfers across corpora and larger |V|, while hyperbolic curvature helps selectively. We also characterize boundary collapse in naive hyperbolic recurrence and show how Mobius stabilization restores trainability. Claims are scoped to controlled small-model comparisons, not full-vocabulary state of the art.