Fast LM Head via Coarse-Fine Decomposition and Tile-Wise Weight Clustering

Abstract

Large language models (LLMs) reuse each generated token as the next input during autoregressive decoding, making per-token latency a critical factor in practical inference performance. In this process, the LM Head performs a large-scale linear projection from the hidden state to the full vocabulary space at every decoding step, becoming one of the major bottlenecks in terms of memory bandwidth and latency. Existing output-layer approximation or vocabulary-pruning approaches can reduce the LM Head cost, but they may cause accuracy degradation due to logit-ranking distortion or task-dependent candidate omission. This paper reinterprets the LM Head as a candidate-screening problem and proposes a coarse-fine LM Head algorithm that restructures it into coarse and fine stages. The proposed method converts the LM Head weights into a compact codebook and an index map through tile-wise k-means clustering, and during online decoding, it computes full-vocabulary coarse scores using tile-wise LUTs and index streaming. Then, exact GEMV or GEMM using the original LM Head weights is performed only for the selected top-<inline-formula><tex-math notation="LaTeX">$k$</tex-math><alternatives><mml:math><mml:mi>k</mml:mi></mml:math><inline-graphic xlink:href="seo-ieq1-3725429.gif"/></alternatives></inline-formula> candidates. Experimental results show that the proposed method limits the average accuracy change to -0.21%p compared with the original LM Head across nine downstream benchmarks, while accelerating the LM Head computation by 5.82<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math><alternatives><mml:math><mml:mo mathvariant="bold">×</mml:mo></mml:math><inline-graphic xlink:href="seo-ieq2-3725429.gif"/></alternatives></inline-formula> on average and up to 12.7<inline-formula><tex-math notation="LaTeX">$\times$</tex-math><alternatives><mml:math><mml:mo>×</mml:mo></mml:math><inline-graphic xlink:href="seo-ieq3-3725429.gif"/></alternatives></inline-formula>. In addition, end-to-end TPOT measurements demonstrate a speed-up of up to 1.18<inline-formula><tex-math notation="LaTeX">$\times$</tex-math><alternatives><mml:math><mml:mo>×</mml:mo></mml:math><inline-graphic xlink:href="seo-ieq4-3725429.gif"/></alternatives></inline-formula> on DGX Spark, and the method operates stably under various decoding lengths, batch sizes, and sampling conditions. The implementation and evaluation scripts are publicly available at <uri>https://github.com/HigherIdeal/CFHead</uri>.