Revisiting SMoE Language Models by Evaluating Inefficiencies with Task Specific Expert Pruning

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Sarkar, Soumajyoti, Lausen, Leonard, Cevher, Volkan, Zha, Sheng, Brox, Thomas, Karypis, George
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866912012901023744
author Sarkar, Soumajyoti
Lausen, Leonard
Cevher, Volkan
Zha, Sheng
Brox, Thomas
Karypis, George
author_facet Sarkar, Soumajyoti
Lausen, Leonard
Cevher, Volkan
Zha, Sheng
Brox, Thomas
Karypis, George
contents Sparse Mixture of Expert (SMoE) models have emerged as a scalable alternative to dense models in language modeling. These models use conditionally activated feedforward subnetworks in transformer blocks, allowing for a separation between total model parameters and per-example computation. However, large token-routed SMoE models face a significant challenge: during inference, the entire model must be used for a sequence or a batch, resulting in high latencies in a distributed setting that offsets the advantages of per-token sparse activation. Our research explores task-specific model pruning to inform decisions about designing SMoE architectures, mainly modulating the choice of expert counts in pretraining. We investigate whether such pruned models offer advantages over smaller SMoE models trained from scratch, when evaluating and comparing them individually on tasks. To that end, we introduce an adaptive task-aware pruning technique UNCURL to reduce the number of experts per MoE layer in an offline manner post-training. Our findings reveal a threshold pruning factor for the reduction that depends on the number of experts used in pretraining, above which, the reduction starts to degrade model performance. These insights contribute to our understanding of model design choices when pretraining with SMoE architectures, particularly useful when considering task-specific inference optimization for later stages.
format Preprint
id arxiv_https___arxiv_org_abs_2409_01483
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Revisiting SMoE Language Models by Evaluating Inefficiencies with Task Specific Expert Pruning
Sarkar, Soumajyoti
Lausen, Leonard
Cevher, Volkan
Zha, Sheng
Brox, Thomas
Karypis, George
Machine Learning
Computation and Language
Sparse Mixture of Expert (SMoE) models have emerged as a scalable alternative to dense models in language modeling. These models use conditionally activated feedforward subnetworks in transformer blocks, allowing for a separation between total model parameters and per-example computation. However, large token-routed SMoE models face a significant challenge: during inference, the entire model must be used for a sequence or a batch, resulting in high latencies in a distributed setting that offsets the advantages of per-token sparse activation. Our research explores task-specific model pruning to inform decisions about designing SMoE architectures, mainly modulating the choice of expert counts in pretraining. We investigate whether such pruned models offer advantages over smaller SMoE models trained from scratch, when evaluating and comparing them individually on tasks. To that end, we introduce an adaptive task-aware pruning technique UNCURL to reduce the number of experts per MoE layer in an offline manner post-training. Our findings reveal a threshold pruning factor for the reduction that depends on the number of experts used in pretraining, above which, the reduction starts to degrade model performance. These insights contribute to our understanding of model design choices when pretraining with SMoE architectures, particularly useful when considering task-specific inference optimization for later stages.
title Revisiting SMoE Language Models by Evaluating Inefficiencies with Task Specific Expert Pruning
topic Machine Learning
Computation and Language
url https://arxiv.org/abs/2409.01483