Skip to main content
Back to timeline
arXivSource publication:

Cross-Layer TopK lets the model split its sparsity budget: lower reconstruction error than per-layer TopK at equal activations, and 2.9× faster training steps

Synopsis

Cross-layer transcoders (CLTs) decompose a large language model's layer activations into interpretable features, but layers differ widely in how hard they are to reconstruct, so giving every layer a fixed number of active latents is a poor fit. Cross-Layer TopK, from ByteDance Seed and collaborators, ranks the latents of all layers together and keeps a fixed total of K per token, letting the model decide how many each layer gets. On Qwen3-8B-Base, at the same 360 activations per token as per-layer TopK, validation FVU falls from 27.64% to 27.07%, and pairing it with sparse decoding cuts training step time to about one third.

Source-provided article image: Cross-Layer TopK: Dynamically Allocating Sparsity Across Layers in Transcoders
arXiv

Interpretation

Sparse selection in cross-layer transcoders moves from a fixed per-layer count to one ranking across all layers, so each token's total activations stay fixed while the split across layers is learned. Per-layer TopK fixed each layer's activation count in advance, yet CLT layers are not equally hard to reconstruct, and intermediate layers are markedly harder. On Qwen3-8B-Base with latent size 32768 and 360 activations per token, validation FVU is 27.07% versus 27.64% for per-layer TopK; in layers 15–19, where per-layer TopK's FVU exceeds 40%, Cross-Layer TopK lowers it by 1.3 to 1.6 points.

The mechanism ranks all layers' latents by one score and keeps only the globally largest K, where K equals the number of layers times the per-layer budget. The score is the rectified pre-activation scaled by the norm of the decoder weights through which that feature writes to later layers, measuring how much output it contributes in total. Training uses MLP input and output activations of Qwen3-8B-Base collected from the Olmo 3 corpus, for 3.2B tokens, with latent size 32768 and a per-token budget K of 360.

A fixed activation count per token lets sparse decoding be applied directly: the encoder is sharded along the feature axis and the decoder along the model dimension, and only the K active latents are exchanged across ranks. JumpReLU's activation count varies per token, so training must monitor L0 and switch between dense and sparse decoding; Cross-Layer TopK fixes the count by construction and needs no monitoring. On 128 GPUs with batch size 8192 and per-layer budget k of 10, the embedding_bag kernel, which fuses gather, scaling and summation, cuts step time from 1.025 s to 0.355 s, reducing compute time by 65% and communication time by 67%.

Perspective

The demonstrated setting is the 36 layers of Qwen3-8B-Base with latent size 32768, a per-token budget of 360 and 3.2B training tokens; circuit scores are averaged over 100 held-out validation prompts, and training speed is measured on 128 GPUs with batch size 8192 and per-layer budget 10. The authors note that compute limits kept them from scaling further.

The authors note that a CLT's decoder grows with the square of the number of layers, so each training run of a large CLT consumes substantial compute, which also kept them from scaling further; the scaling-law exponents are of the same order as those reported for TopK SAEs over a comparable range, but were fit only for latent sizes from 213 to 215 and per-layer budgets from 5 to 20.