Model Compression Method for S4 with Diagonal State Space Layers using Balanced Truncation
This work addresses model compression for S4 models on edge devices, offering an incremental improvement by adapting an existing method to a specific architecture.
The paper tackles the problem of compressing S4 models with Diagonal State Space layers for edge deployment by applying balanced truncation, a control theory technique, and using the reduced parameters as initialization, achieving higher accuracy with fewer parameters compared to conventional training.
To implement deep learning models on edge devices, model compression methods have been widely recognized as useful. However, it remains unclear which model compression methods are effective for Structured State Space Sequence (S4) models incorporating Diagonal State Space (DSS) layers, tailored for processing long-sequence data. In this paper, we propose to use the balanced truncation, a prevalent model reduction technique in control theory, applied specifically to DSS layers in pre-trained S4 model as a novel model compression method. Moreover, we propose using the reduced model parameters obtained by the balanced truncation as initial parameters of S4 models with DSS layers during the main training process. Numerical experiments demonstrate that our trained models combined with the balanced truncation surpass conventionally trained models with Skew-HiPPO initialization in accuracy, even with fewer parameters. Furthermore, our observations reveal a positive correlation: higher accuracy in the original model consistently leads to increased accuracy in models trained using our model compression method, suggesting that our approach effectively leverages the strengths of the original model.