UNAST: Unified framework for Neural Architecture Search for Transformers

Ilia Markov, Chenhan D. Yu, Hongxu Yin, Saurav Muralidharan, Greg Heinrich, Jan Kautz, Dan Alistarh, Pavlo Molchanov

Keywords: Efficiency, Neural Architecture Search, Large Language models

Abstract: We introduce the UNAST, a new approach to optimize Large Language Models (LLMs) post-training. UNAST combines Neural Architecture Search (NAS) with sparsity and quantization for LLM compression. Starting with a trained model, UNAST replaces layers (e.g., attention and MLP) with more efficient alternatives by adjusting attention heads, KV projection dimensions, and MLP expansion factors. Local distillation pretrains layer candidates to mimic original layers. Scores and costs (latency, number of parameters, etc.) of each operator are fed into an Integer Linear Optimizer to find the optimal architecture under predefined constraints (latency, number of parameters, etc.). Our experiments show that UNAST scales to large models, reducing training costs by up to 10 times compared to training smaller models from scratch. Validation on GPT-3 and LLaMa models demonstrate that UNAST improves latency and memory footprint by up to 60\% with minimal accuracy loss. UNAST also provides insights into the effects of different compression types on Transformer layers, aiding in the development of non-uniform models.

Primary Area: foundation or frontier models, including LLMs

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6615