Tuning Language Models by Mixture-of-Depths Ensemble

TL;DR

This paper introduces Mixture-of-Depths (MoD), a novel framework that leverages intermediate layers of transformer-based language models through learned ensemble routing, improving performance and efficiency in language modeling tasks.

Contribution

The paper proposes MoD, a new tuning method that utilizes intermediate layers as ensembles with learned routing, enhancing language model training and reducing parameters needed.

Findings

MoD improves performance across various language modeling tasks.

Using intermediate layers as ensembles can match or surpass final-layer tuning.

MoD achieves similar results with fewer trainable parameters.

Abstract

Transformer-based Large Language Models (LLMs) traditionally rely on final-layer loss for training and final-layer representations for predictions, potentially overlooking the predictive power embedded in intermediate layers. Surprisingly, we find that focusing training efforts on these intermediate layers can yield training losses comparable to those of final layers, with complementary test-time performance. We introduce a novel tuning framework, Mixture-of-Depths (MoD), which trains late layers as ensembles contributing to the final logits through learned routing weights. With the auxiliary distillation loss and additional normalization modules, we ensure that the outputs of the late layers adapt to language modeling. Our MoD framework, which can be integrated with any existing tuning method, shows consistent improvement on various language modelling tasks. Furthermore, by replacing traditional trainable modules with MoD, our approach achieves similar performance with significantly fewer trainable parameters, demonstrating the potential of leveraging predictive power from intermediate representations during training.