Technical Summary: Leviathan – Decoupling Input and Output Representations in Language Models

1. Problem Statement

Modern large language models (LLMs) typically employ a tied-weight architecture, where a single embedding matrix $E \in \mathbb{R}^{V \times D}$ serves dual roles: mapping input tokens to continuous representations and projecting hidden states back into vocabulary space for output prediction. This coupling forces a single matrix to simultaneously satisfy two distinct objectives:

1. Input Representation: Requires smoothness, where semantically related tokens occupy nearby regions in the representation space.
2. Output Classification: Requires discrimination, effectively separating all tokens across the vocabulary.

The authors hypothesize that this coupling imposes an architectural limitation. While fully untied models (learning separate input and output matrices) theoretically decouple these objectives, they prove ineffective in practice. Experiments show that fully untied models underperform tied baselines despite a ~50% increase in parameter count, suggesting that the additional parameters are not utilized effectively due to optimization difficulties and insufficient training signal for discrete embeddings, particularly for rare tokens.

2. Methodology: The Leviathan Architecture

Leviathan introduces a Transformer architecture that decouples input representation from output prediction without incurring the prohibitive parameter cost of fully untied embeddings.

2.1 Learned Embedding Vectorization (LEV)

The core innovation is the LEV layer, which replaces the dense input embedding matrix with a learned continuous function $G: \{0, \dots, V-1\} \to \mathbb{R}^D$. This function maps token indices to embeddings through a compact, parameter-sharing mechanism:

• Compositional Indexing: The vocabulary is factorized into $k$ components using a base-$b$ decomposition, converting a token index $i$ into a tuple $(i_1, \dots, i_k)$.
• Shared Codebooks: These indices map to shared codebooks $C_1, \dots, C_k$, yielding a seed representation $z(i)$. This allows rare tokens to benefit from gradient signals passed through shared entries used by structurally similar tokens.
• Continuous Projection: The seed is projected and normalized to a bounded latent coordinate $\tilde{z}_\ell \in [0, 1]^{d_{seed}}$.
• Spline Expansion: Each dimension is modeled by a univariate quadratic B-spline basis expansion with $\kappa$ knots.
• Tensor Product: Dimensions are combined via a rank-$r$ tensor product.
• Head Summation: The final embedding is the sum of outputs from $h$ parallel heads after a linear projection.

2.2 Architecture and Decoupling

• Untied Output Head: The LEV layer allows the output projection head to remain fully untied (a dense matrix independent of the input).
• Parameter Efficiency: Because the LEV layer is significantly smaller than a dense embedding table, the overhead of the untied output head is negligible. At the 1.2B parameter scale, the total parameter increase is only 0.2%.
• Inductive Bias: The architecture imposes a smoothness inductive bias over the token space, contrasting with the discrete, potentially anisotropic geometry of standard lookup tables.

3. Experimental Setup

• Models: Experiments were conducted on decoder-only Transformers at three scales: 200M, 400M, and 1.2B parameters.
• Baselines:
  • TIED: Standard tied-weight baseline.
  • UNTIED: Fully untied model (611M parameters at 400M scale).
  • ALBERT: Input-parameter-matched factorization (linear compression) with an untied output head.
  • Ablations: Parameter-matched Flat MLP and Factored MLP to isolate the contributions of nonlinearity and compositional structure.
• Training: All models were trained on The Pile (uncopyrighted) with identical hyperparameters, seeds, and token budgets (up to 24.1B tokens for 1.2B models).
• Evaluation: Language modeling performance on a held-out validation set of The Pile, zero-shot generalization on Wikitext-103, and downstream benchmarks (HellaSwag, ARC, PIQA, WinoGrande, LAMBADA).

4. Key Results

4.1 Language Modeling Performance

Leviathan consistently outperforms the tied-weight baseline across all scales, with the performance gap widening as training progresses and model scale increases.

• 1.2B Scale: Leviathan reduced validation perplexity by 9.2% (0.087 nats) compared to the tied baseline.
• Sample Efficiency: At 1.2B scale, Leviathan required 2.1× fewer training tokens to reach the tied baseline's final loss.
• Training Dynamics: At smaller scales (200M), Leviathan initially underperformed before crossing the baseline; at 1.2B, it outperformed the baseline throughout training.

4.2 Comparison to Alternatives

• Vs. Fully Untied: The fully untied model (UNTIED) underperformed the tied baseline throughout training, confirming that naive decoupling without efficient parameterization is suboptimal.
• Vs. Linear Factorization: The ALBERT-style factorized baseline also underperformed the tied baseline, indicating that linear compression alone cannot recover the representational benefits of the tied geometry.
• Vs. Leviathan: Leviathan outperformed both UNTIED and ALBERT, demonstrating that nonlinear continuous generation is superior to both naive untying and linear factorization within Chinchilla-optimal compute budgets.

4.3 Downstream and Generalization

• Downstream Tasks: At 1.2B scale, Leviathan improved performance on all six evaluated downstream benchmarks. Notably, it achieved a 30% reduction in perplexity on LAMBADA.
• Zero-Shot Transfer: On Wikitext-103 (a different distribution), Leviathan showed increasing perplexity improvements with scale (2.4% at 200M, 12.3% at 1.2B), suggesting its continuous representations transfer more effectively across data distributions.

4.4 Frequency-Stratified Analysis

The performance gains are heavily concentrated in the long tail of the token distribution:

• Rare Tokens: For tokens appearing fewer than 1-in-10 million times, Leviathan reduced perplexity by 81% at 1.2B scale.
• Frequent Tokens: For tokens appearing more than 1-in-10 thousand times, the improvement dropped to near zero (and slightly negative for the most frequent), as discrete embeddings are sufficient when training signals are strong.
• Mechanism: The continuous parameterization shares gradients across structurally similar tokens, providing a robust signal where discrete lookups fail due to data sparsity.

4.5 Robustness

• Vocabulary Size: Results remained consistent across vocabulary sizes of 50k, 100k, and 200k, with relative improvements of ~2% across the board.
• Ablations: The gains were decomposed into: continuous representation (+1.13%), compositional structure (+0.37%), and separable spline smoothness (+0.62%).

5. Significance and Claims

The paper claims that the standard practice of weight-tying imposes an architectural limitation that is not resolved by simply untying weights. Instead, the optimal solution involves decoupling input and output representations while enforcing smoothness via continuous parameterization.

• Architectural Insight: The representational geometry of input embeddings matters. Continuous parameterization is a data- and compute-efficient means of improving this geometry, particularly for rare tokens where discrete embeddings struggle.
• Efficiency: Leviathan achieves these gains with a negligible parameter overhead (0.2% at 1.2B scale), making it a practical upgrade for modern language models.
• Scalability: The benefits of the LEV layer compound with scale and data, suggesting that as models grow, the advantage of smooth, continuous input representations over discrete lookups becomes more pronounced.

The authors conclude that while the frequency-dependent improvement pattern is correlational, the consistency across scales, vocabularies, and architectural variants provides robust evidence that continuous parameterization effectively mitigates the optimization difficulties associated with learning discrete embeddings for rare tokens.