Amazon

State Space Model Optimization for Edge AI

State Space Models like Mamba 2 attention compute scales linearly instead of quadratic, which is a benefit. Another benefit is that that KV cache is not scaling linearly but remains constant. Amazon Devices and Services recognizes this benefit as required RAM and compute reduces for such models. However, in addition to the benefits, there is scientific risk that when these models are optimized or quantized they don't maintain accuracy. Thus, this student team will work to invent scientific techniques to quantize and optimize SSM and create multiple optimal models on pareto optimal curve with accuracy and performance. The second part of the project this student team will work to achieve is to run these models accuracy once optimized on edge hardware like Orange Pi or Raspberry Pi using open source frameworks like Ollama, Llama.cpp. Research paper: https://arxiv.org/pdf/2405.21060 Accuracy: Standard benchmarks as defined in LM Harness like Lambada, ARC easy/challenge. Performance Criterion: Accuracy degradation after model optimization and quantization <3% degradation on each task Demonstration: Optimized model running on Edge HW like Orange Pi or Raspberry Pi The outcomes this student team will work to achieve are: Quantized Mamba 2 or equivalent LLM architecture with <3% accuracy degradation on LM Harness Quantized model deployed on Orange Pi using open source frameworks like Ollama, Llama.cpp The bonus deliverables this student team will work to achieve are: Quantized + Sparse (50%) on Mamba 2 or equivalent LLM architecture with <3% accuracy degradation on LM Harness Quantized model deployed on Orange Pi with access of NPU from Rochip on Orange Pi