PaperlayerView exact revision on arXiv

Source-linked AI summary

RiLM: Parameter-Efficient Language Modeling via Geodesic Decoding

Fang Li

arXiv:2609.10305v1cs.CLcs.AI

TL;DR

Small language models can devote a large share of their parameters to vocabulary-sized output layers, motivating a geometry-based alternative. RiLM removes that layer by evolving context on a manifold and decoding through distances to vocabulary embeddings; in controlled comparisons, HypRiLM is strongest on WikiText-2, while stable Möbius recurrence is necessary for hyperbolic training.

  • Problem

    At sub-million scale, the vocabulary-sized output matrix can consume roughly one third of model capacity, while Euclidean decoding does not directly reflect how states move through representation space.

  • Method

    RiLM uses a shared composition map to move context along manifold geodesics and derives next-token probabilities from squared geodesic distances to vocabulary embeddings.

  • Results

    54.2±0.2 validation perplexity on WikiText-2 makes HypRiLM stronger than Flat RiLM at 87.6±0.6 and the reported tied and matched recurrent and attention controls.

  • Takeaways & Limitations

    Geometry can replace the output layer in controlled small-model language modeling, with hyperbolic curvature helping on WikiText-2 but not uniformly across settings.

  • Takeaways & Limitations

    The study is bounded to controlled settings with fewer than 2M parameters, context 64, TBPTT k=8, and a 2k vocabulary that is not a full WikiText-2 benchmark.

Abstract

from arXiv · show

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.

I. INTRODUCTION

RiLM replaces the vocabulary-sized output layer with geodesic decoding, making context a manifold trajectory whose vocabulary embeddings serve both input and output. The paper evaluates flat and hyperbolic variants under controlled small-model comparisons, emphasizing parameter efficiency and trainability.

  • Motivation: W_out can consume one third of a sub-million-parameter model at d=128 and |V|=2000, motivating alternatives to conventional linear decoding.Tied embeddings reduce redundant rows but still decode through a linear functional.
  • RiLM: RiLM evolves context on a Riemannian manifold with a shared composition map and predicts words from squared geodesic distance to vocabulary embeddings.This couples representation and classification without a separate W_out layer.
  • RiLM: Flat RiLM uses R^d, whereas HypRiLM uses the Poincaré ball H^d to test whether curvature improves sequential composition with the same φ.The comparison isolates geometry while retaining the shared template.
  • Trainability: Naive hyperbolic recurrence causes boundary collapse, while Möbius-based updates restore stable training and usable hyperbolic dynamics.The failure drives logits toward uniform prediction; the stabilization is presented as necessary for the hyperbolic model.
  • Evaluation scope: Both RiLM variants outperform tied and matched LSTM, Transformer, and SSM controls on WikiText-2 and Penn Treebank under the paper’s fairness protocols.The study does not target billion-parameter pretraining or full-vocabulary leaderboard performance.

III. METHOD

RiLM treats sentences as manifold trajectories: tokens update a shared state, and decoding measures geometric proximity between that state and vocabulary embeddings. Hyperbolic space supplies curvature and boundary-sensitive geometry for this construction.

  • Geodesic decoding: RiLM replaces a separate classifier with negative squared geodesic distance between the contextual state and each vocabulary embedding.The design mirrors tied embeddings while changing the readout from a linear map to geometry.
  • Manifold operations: A Riemannian manifold provides tangent-space inner products, exponential and logarithmic maps, and shortest-path geodesic distance.In Euclidean space, the exponential map is addition and geodesic distance is the usual norm.
  • Hyperbolic geometry: The Poincaré ball models hyperbolic space with curvature c>0 and exponentially growing volume, while its conformal factor diverges at the boundary.Hyperbolic distance is computed through Möbius displacement within the open ball.

B. Recurrence and Geodesic Decoding

RiLM composes token inputs through a shared tangent-space map, advances the state along a geodesic, and converts distances to vocabulary embeddings into next-token probabilities. In Euclidean space, this becomes additive recurrence.

  • Recurrence: The shared composition map φ reads the current state and incoming token in tangent coordinates and proposes the update direction.The contextual state and vocabulary embeddings occupy the same manifold.
  • Decoding: Temperature-scaled decoding applies a softmax over negative squared geodesic distances, favoring vocabulary embeddings near the current state.This makes likely continuations those inside a small geodesic neighborhood.
  • Euclidean special case: On R^d, the geodesic recurrence reduces to h_t+1 = h_t + φ(h_t, e_wt), the Euclidean special case of RiLM.The sequence begins from a learned root state.
  • Recurrence and decoding: Figure 1 summarizes one timestep as shared composition, geodesic state updating, and distance-based decoding without W_out.The same embedding map supports token ingestion and prediction.

C. Hyperbolic RiLM and M¨obius Recurrence

Direct exponential-map updates in hyperbolic RiLM drive states toward the ball boundary and flatten distances, preventing learning. Möbius transport computes increments in a well-behaved region and restores stable recurrence.

  • Boundary collapse: As hyperbolic states approach the boundary, the conformal factor diverges, making small tangent steps produce enormous motion and nearly constant vocabulary distances.The observed result is boundary sticking and perplexity equal to the vocabulary size.
  • Möbius recurrence: Möbius stabilization computes increments at the origin and transports them with Möbius addition, keeping states inside the ball.This avoids the direct exponential-map pathology.
  • Stability: After training, held-out states typically remain at norms 0.29–0.71, whereas projecting the state after every step stalls optimization near 130 PPL.The implementation projects embeddings and the root parameter instead of every recurrent state.

D. Composition Function and Training

RiLM uses one timestep-shared MLP composition map and trains with masked token likelihood under truncated backpropagation. The main experiments use context length 64 and compare shared training settings across models.

  • Composition Function: φ(vh, ve) = tanh(W[vh; ve] + b) uses W ∈ R^d×2d and is shared across timesteps.At d=128, φ has approximately 33k parameters.
  • Training Objective: The model minimizes masked token negative log-likelihood for next-token prediction.
  • Training Procedure: Truncated BPTT with k=8 detaches states every eight steps, limiting gradient depth through repeated φ.Each token participates in eight compositional steps before gradients truncate.
  • Shared Settings: All models train for 10 epochs with batch size 128 and context length 64.Flat RiLM uses η = 10^-3, while HypRiLM uses η = 3 × 10^-3, curvature c = 1.0, and gradient clipping at norm 1.0.

E. Parameter Budget and Baseline Fairness

The evaluation separates decoding effects from parameter-count effects by comparing tied models at d=128 with matched models near 290k parameters. Validation perplexity is measured on controlled 2k-vocabulary WT-2 and PTB setups rather than full-vocabulary benchmarks.

  • Parameter Budget: Geodesic decoding removes the approximately 256k-parameter Wout at d=128 and |V|=2000, while RiLM retains full-width embeddings and φ.Tied baselines share input/output embeddings but retain flat LSTM, attention, or SSM dynamics.
  • Baseline Fairness: Matched baselines equalize total parameters at approximately 290k by shrinking hidden width to d′ ≈82–97.This regime tests whether RiLM benefits only from retaining d=128.
  • Evaluation Protocol: The study reports validation perplexity on WikiText-2 and Penn Treebank using the 2000 most frequent words.The fixed vocabulary makes Wout's parameter cost explicit, but absolute perplexities are not comparable to full-vocabulary leaderboard entries.
  • Baseline Fairness: Baselines include two-layer LSTM, causal Transformer, and selective diagonal SSM models in untied, tied, and matched regimes.The SSM selects its best validation checkpoint, whereas other models use the final epoch checkpoint.

A. WikiText-2 at Controlled Vocabulary

Across controlled comparisons, RiLM variants outperform the evaluated tied and matched baselines on WT-2 and PTB, while curvature is dataset-dependent. Geodesic decoding remains advantageous at 10k vocabulary, but absolute perplexities are not comparable across vocabulary settings.

  • WT-2 Results: 54.2±0.2 perplexity is achieved by HypRiLM on WT-2, versus 87.6±0.6 for Flat RiLM at approximately 289k parameters.Hyperbolic space provides an approximately 38% relative improvement on this corpus.
  • WT-2 Results: 113.0±3.8 perplexity is reported for SSM-tied, while HypRiLM reaches 54.2 on WT-2.HypRiLM leads tied LSTM, Transformer-tied, and SSM-tied controls by 59–93 PPL.
  • WT-2 Results: 64–89 PPL separates HypRiLM from matched approximately 290k baselines, with no fair regime reversing the ranking.Matched baselines shrink hidden widths to d′ ≈82–97 while HypRiLM retains d=128.
  • Stability: Ablating Möbius stabilization drives perplexity to the uniform |V|=2000 baseline, showing that stable hyperbolic dynamics are prerequisite to the headline results.
  • Vocabulary Scaling: The 10k-vocabulary setting makes absolute perplexity incomparable to Table III because embedding tables dominate parameter count.The comparison therefore tests whether model ordering persists rather than absolute performance.
  • Vocabulary Scaling: At 10k vocabulary, Flat RiLM reaches 341.8±0.5 PPL and HypRiLM 345.8±16.0, both below SSM-tied at 708.3±9.0.The ranking persists, but Flat RiLM is marginally best and HypRiLM has an outlier seed.
  • Cross-Corpus Behavior: 40.9±0.6 PTB perplexity for Flat RiLM beats 69.8±0.5 for HypRiLM at identical approximately 289k parameters.Both variants remain below tied and matched LSTM, Transformer, and SSM controls.

D. Ablations and Literature Context

The ablations attribute HypRiLM’s WikiText-2 advantage primarily to hyperbolic geometry rather than composition-map choice, while literature comparisons require caution because vocabulary and scale differ. Geometry analyses further show stable interior trajectories and plausible distance-ranked continuations.

  • Ablations: A ∼34 PPL gap separates hyperbolic and flat MLP rows, while spline composition changes Flat RiLM by ∼4.4 PPL but HypRiLM by only ∼0.1 PPL.Reducing curvature to c=0.1 also degrades HypRiLM from 54.0 to 71.1 PPL.
  • Literature context: HypRiLM uses ∼289k parameters and a 2k-word vocabulary, so its 54.0 PPL is not a controlled comparison with full-vocabulary 33M–257M-parameter literature models.The reported proximity to AWD-LSTM’s 60.7 PPL is presented only as context, not competitive state of the art.
  • Learned geometry: HypRiLM trajectories remain inside the Poincaré ball, with ∥h_t∥∈[0.29, 0.71] on the examined validation prefix, while Flat RiLM follows a comparable low-dimensional path in R^d.The geometry inspection uses WT-2, seed 42, ten epochs, and a held-out validation prefix.
  • Distance-based decoding: After “along with · · · city of,” both models rank “was” and related function words among the highest-probability continuations.The example is the first long validation prefix returned by the data loader rather than a cherry-picked case.

V. DISCUSSION

Naive exponential-map recurrence collapses hyperbolic states at the boundary, flattening logits and preventing learning. Möbius stabilization keeps states interior and restores usable training dynamics.

  • Boundary collapse: Within three recurrence steps, naive exp_h_t updates reach ∥h_t∥≈0.999, equalize vocabulary distances, flatten logits, and produce PPL = |V| = 2000.This boundary collapse makes naive hyperbolic recurrence untrainable for language modeling at modest depth.
  • Stabilization: Möbius updates compute tangent increments at the origin and translate them with ⊕_c, keeping states interior; held-out trajectories remain within ∥h_t∥∈[0.29, 0.71].Hard projection after every step instead yields ∼130 PPL, so the method projects only embeddings and θ_root.

B. Why Geodesic Decoding Helps at Small Scale

RiLM makes vocabulary embeddings serve both input and output, replacing a vocabulary-sized linear readout with metric decoding. Hyperbolic geometry helps substantially on WT-2 in these experiments but is not uniformly preferable across corpora.

  • Parameter efficiency: When |V| and d are moderate, W_out consumes roughly one third of an untied LM’s parameters; RiLM instead retains d=128 and reuses the vocabulary table for distance-based classification.This avoids shrinking hidden width to match parameter budgets.
  • When curvature helps: Swapping R^d for H^d at fixed φ yields a large WT-2 gain but not a PTB gain, so manifold choice should be made using validation performance.The paper associates the WT-2 pattern with broader topical and lexical structure, while compact PTB sentences are adequately served by flat geometry.
  • Operational decoding: The WT-2 example after “along with · · · city of” ranks plausible function-word continuations using logits defined as −d_M^2/τ.The table reports top-5 next-token candidates by geodesic distance.
  • Relation to tied embeddings: Geodesic decoding differs from tied linear decoding because word-dependent embedding norms and metric distances shape the logits, not only inner products.The tied baselines remove W_out but retain linear scoring and still trail HypRiLM by 59+ PPL on WT-2.

D. Efficiency and Limitations

RiLM reduces parameter cost by replacing the output matrix with distance-based decoding, while efficiency depends on geometry and vocabulary size. The study remains bounded to controlled small-model settings, with scaling and timing limitations.

  • Efficiency: 0.15 ms/token for Flat RiLM, versus 0.52 ms/token for HypRiLM, 0.44 ms/token for LSTM, and 0.40 ms/token for Transformer on WikiText-2.Hyperbolic operations add approximately 3.5× overhead over Flat RiLM, while Flat RiLM is the fastest variant in this benchmark.
  • Efficiency: O(|V| · d) distance evaluation has the same asymptotic order as softmax, so RiLM’s benefit is parametric rather than computational at small vocabularies.The approach removes W_out but does not eliminate vocabulary-dependent decoding work.
  • Limitations: RiLM’s evaluation is bounded to fewer than 2M parameters, context 64, TBPTT k=8, and controlled 2k- or 10k-vocabulary settings.The 2k vocabulary is explicitly described as a controlled instrument rather than a full WikiText-2 benchmark.
  • Limitations: At full vocabulary, RiLM retains O(|V| · d) decoding cost, and geodesic distance over the full table may require approximation.The paper identifies hierarchical softmax, sampled negatives, or product-manifold factorization as possible directions, while leaving curvature benefits at scale unresolved.
  • Takeaway: Within the controlled regime, geodesic decoding couples representation and prediction under parameter scarcity but is not presented as a replacement for billion-parameter pretraining.The reported benefits persist under tied and matched comparisons and partially under larger vocabularies.
Loading 2609.10305v1…