Layer-Adaptive State Pruning for Deep State Space Models

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Gwak, Minseon, Moon, Seongrok, Ko, Joohwan, Park, PooGyeon
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915130553401344
author Gwak, Minseon
Moon, Seongrok
Ko, Joohwan
Park, PooGyeon
author_facet Gwak, Minseon
Moon, Seongrok
Ko, Joohwan
Park, PooGyeon
contents Due to the lack of state dimension optimization methods, deep state space models (SSMs) have sacrificed model capacity, training search space, or stability to alleviate computational costs caused by high state dimensions. In this work, we provide a structured pruning method for SSMs, Layer-Adaptive STate pruning (LAST), which reduces the state dimension of each layer in minimizing model-level output energy loss by extending modal truncation for a single system. LAST scores are evaluated using the $\mathcal{H}_{\infty}$ norms of subsystems and layer-wise energy normalization. The scores serve as global pruning criteria, enabling cross-layer comparison of states and layer-adaptive pruning. Across various sequence benchmarks, LAST optimizes previous SSMs, revealing the redundancy and compressibility of their state spaces. Notably, we demonstrate that, on average, pruning 33% of states still maintains performance with 0.52% accuracy loss in multi-input multi-output SSMs without retraining. Code is available at https://github.com/msgwak/LAST.
format Preprint
id arxiv_https___arxiv_org_abs_2411_02824
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Layer-Adaptive State Pruning for Deep State Space Models
Gwak, Minseon
Moon, Seongrok
Ko, Joohwan
Park, PooGyeon
Machine Learning
Systems and Control
Due to the lack of state dimension optimization methods, deep state space models (SSMs) have sacrificed model capacity, training search space, or stability to alleviate computational costs caused by high state dimensions. In this work, we provide a structured pruning method for SSMs, Layer-Adaptive STate pruning (LAST), which reduces the state dimension of each layer in minimizing model-level output energy loss by extending modal truncation for a single system. LAST scores are evaluated using the $\mathcal{H}_{\infty}$ norms of subsystems and layer-wise energy normalization. The scores serve as global pruning criteria, enabling cross-layer comparison of states and layer-adaptive pruning. Across various sequence benchmarks, LAST optimizes previous SSMs, revealing the redundancy and compressibility of their state spaces. Notably, we demonstrate that, on average, pruning 33% of states still maintains performance with 0.52% accuracy loss in multi-input multi-output SSMs without retraining. Code is available at https://github.com/msgwak/LAST.
title Layer-Adaptive State Pruning for Deep State Space Models
topic Machine Learning
Systems and Control
url https://arxiv.org/abs/2411.02824