Research map / Architectures and efficiency / Attention and transformers

Leviathan: Decoupling Input and Output Representations in Language Models

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Language model pretraining and induction · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Replaces the input embedding matrix of a decoder-only Transformer with a compact learned continuous mapping (LEV layer, base-b factorized indices, B-spline basis) while leaving the output head untied, and trains 200M–1.2B models on The Pile against tied, untied, and ALBERT-style baselines.

Problem

Weight tying couples the input embedding and output classifier roles in language models, and naively untying embeddings adds many parameters without improving performance, so a better input–output parameterization is needed.

Approach

Input embeddings are generated by a learned continuous function over compositionally factorized token indices (shared codebooks, per-head latent projection, separable quadratic B-spline expansion, rank-r tensor product), which is cheap enough to allow a fully untied output head with under 0.2% extra parameters at 1.2B.

Key results

• At 1.2B, Leviathan reaches 2.402 nats vs 2.489 for Tied (9.2% perplexity reduction) and needs 2.1x fewer tokens to reach the tied baseline's final loss.
• At 1.2B, Leviathan improves on all six downstream benchmarks, with a 30% LAMBADA perplexity reduction relative to Tied.
• At 400M, Leviathan (2.724 nats) beats Tied (2.784), Untied (2.801) and ALBERT-style factorization (2.807), with the gain concentrated in rare tokens (81% perplexity reduction for tokens seen 1-9 times at 1.2B).

Surprising findings

• Fully untied embeddings underperform the tied baseline throughout training despite about 50% more parameters at 400M.
• Linear low-rank factorization of the input embedding (ALBERT-style) is worse than the tied baseline, while the nonlinear continuous LEV layer is better.

Limitations

• Single training corpus (The Pile), single tokenizer family (BPE), and maximum scale of 1.2B parameters.
• Each configuration trained with a single seed, so no formal confidence intervals per scale.
• Frequency-dependent analysis is correlational; behavior at 7B+ scale and multilingual corpora is untested.

Experimental setup

Datasets: The Pile WikiText-103

Benchmarks: The Pile validation WikiText-103 HellaSwag ARC-Easy ARC-Challenge PIQA WinoGrande LAMBADA

Metrics: validation loss perplexity zero-shot accuracy normalized accuracy per-frequency-bucket NLL

Baselines: Tied embeddings Untied embeddings ALBERT-style factorized embeddings Flat MLP Factored MLP

Builds on: weight tying ALBERT factorized embeddings rotary positional embeddings SwiGLU Chinchilla scaling

Closest papers by meaning

• DEPT: Decoupled Embeddings for Pre-training Language Models (ICLR 2025 · similarity 0.90)
• Scaling Embedding Layers in Language Models (NeurIPS 2025 · similarity 0.89)
• Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling (ICML 2025 · similarity 0.88)
• Language Models as Hierarchy Encoders (NeurIPS 2024 · similarity 0.88)
• Transformer Block Coupling and its Correlation with Generalization in LLMs (ICLR 2025 · similarity 0.88)
• Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch (ICML 2024 · similarity 0.88)