Leviathan: Decoupling Input and Output Representations in Language Models

TL;DR

Leviathan introduces a novel Transformer architecture that decouples input and output representations, leading to improved language modeling performance, especially on rare tokens, with minimal additional parameters.

Contribution

The paper presents Leviathan, a new method replacing input embeddings with learned vectorization, enhancing performance over standard tied embeddings with minimal parameter increase.

Findings

Leviathan reduces validation perplexity by 9% at 1.2B scale.

It requires 2.1 times fewer tokens to reach baseline loss.

Achieves a 30% reduction in LAMBADA perplexity.

Abstract

Modern language models use a single matrix for input embedding and output projection. This couples two distinct objectives: token representation and discrimination over a vocabulary. This work introduces Leviathan, a Transformer architecture that replaces the input embedding matrix with learned embedding vectorization (LEV), a compact continuous mapping from token indices to embeddings. Leviathan's output head remains untied for a parameter increase of as low as 0.2%. Under controlled comparisons with identical Transformer backbones, Leviathan consistently improves language modeling performance over standard tied-embedding baselines across a 200M-1.2B parameter regime on The Pile with gains that grow during training. At 1.2B scale, Leviathan reduces validation perplexity by 9%, requires $2.1\times$ fewer training tokens to reach the tied baseline's final loss, and improves on all six downstream benchmarks evaluated, including a 30% reduction in LAMBADA perplexity. Frequency-stratified analysis reveals gains to be concentrated in rare tokens, where continuous parameterization reduces perplexity by 81%, falling to near zero for the most frequent.