Balcony: A Lightweight Approach to Dynamic Inference of Generative Language Models

TL;DR

Balcony addresses the challenge of deploying large generative LLMs under strict computational constraints by freezing the base model and inserting a minimal, single-transformer Balcony layer at selected exit points. It trains these Balcony layers with a self-distillation objective based on the $KL$ divergence between the full-model and sub-model outputs, using a shared LM head to maintain consistency across exits. The approach yields significant speedups with minimal performance loss, outperforming state-of-the-art dynamic-inference methods such as Flextron and LayerSkip across multiple models and benchmarks while dramatically reducing training data and parameter requirements. This simplicity and efficiency enable practical, on-the-fly adaptation to varying budgets and hardware contexts, with potential for further improvements via integration with MoE, Mamba, or speculative decoding techniques.

Abstract

Deploying large language models (LLMs) in real-world applications is often hindered by strict computational and latency constraints. While dynamic inference offers the flexibility to adjust model behavior based on varying resource budgets, existing methods are frequently limited by hardware inefficiencies or performance degradation. In this paper, we introduce Balcony, a simple yet highly effective framework for depth-based dynamic inference. By freezing the pretrained LLM and inserting additional transformer layers at selected exit points, Balcony maintains the full model's performance while enabling real-time adaptation to different computational budgets. These additional layers are trained using a straightforward self-distillation loss, aligning the sub-model outputs with those of the full model. This approach requires significantly fewer training tokens and tunable parameters, drastically reducing computational costs compared to prior methods. When applied to the LLaMA3-8B model, using only 0.2% of the original pretraining data, Balcony achieves minimal performance degradation while enabling significant speedups. Remarkably, we show that Balcony outperforms state-of-the-art methods such as Flextron and Layerskip as well as other leading compression techniques on multiple models and at various scales, across a variety of benchmarks.

Figures (8)

• Figure 1: An overview of the Balcony inference. Balcony preserves the base model’s performance while enabling efficient, on-the-fly adaptation to different computational budgets.
• Figure 2: Speed-up as a function of the ratio of active parameters in modified versions of Llama 3 8B. Comparisons are shown between width-only and depth-only modifications, with the unmodified baseline included for reference.
• Figure 3: Training in the Balcony framework: By freezing the pretrained base LLM and adding a decoder layer at each exit point, Balcony can outperform SoTA with significantly fewer training tokens.
• Figure 4: The Balcony-LLaMA2-7B model family demonstrates superior performance on MMLU and ARC-E compared to Flextron-Dynamic/Static, LayerSkip and post-hoc compression methods like Compresso, LLM-Pruner, SliceGPT, and LaCo.
• Figure 5: Ablation studies on the Balcony framework. From left to right: (1) Impact of random initialization and the contributions of MLP-only and Attention-only modules. (2) Effect of incorporating Cross-Entropy (CE) loss in self-distillation. (3) Effect of freezing the main architecture during post-training. The results are the average scores across ARC-E, ARC-C, BoolQ, HellaSwag, Lambada, MMLU, OpenBookQA, PIQA, and Winogrande.