🤖 machine learning

Flash PD-SSM: Memory-Optimized Structured Sparse State-Space Models

본 논문은 대규모 학습에 필요한 계산 효율성을 유지하면서 구조화되지 않은 모델의 높은 표현력을 달성하고 장문 시퀀스에서의 최첨단 성능을 보장하기 위해 학습 가능한 희소 행렬 집합에서 동적으로 선택하는 메모리 최적화 구조화 희소 상태 공간 모델인 Flash PD-SSM 을 소개합니다.

원저자: Aleksandar Terzić, Francesco Carzaniga, Nicolas Menet, Yannick Biehl, Michael Hersche, Thomas Hofmann, Abbas Rahimi

게시일 2026-05-20

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Aleksandar Terzić, Francesco Carzaniga, Nicolas Menet, Yannick Biehl, Michael Hersche, Thomas Hofmann, Abbas Rahimi

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

상상해 보세요. 긴 이야기를 읽고, 세부 사항을 기억하며, 다음에 무슨 일이 일어날지 예측할 수 있는 초지능 로봇을 만들려고 합니다. 이를 위해 로봇은 읽는 동안 정보를 유지할 수 있는 '뇌'가 필요합니다.

오랜 시간 동안 이 작업에 가장 적합한 '뇌'는 트랜스포머(현재 많은 AI 채팅봇의 기반 기술) 와 같았습니다. 이들은 놀라울 정도로 똑똑하지만 매우 무겁고 느립니다. 마치 엄청난 양의 연료를 소비하고 거대한 차고가 필요한 고급 리무진과 같습니다.

그런 다음 엔지니어들은 **상태 공간 모델 **(SSM)을 발명했습니다. 이것들을 전기 킥보드라고 생각하세요. 훨씬 더 빠르고 에너지 (메모리) 를 훨씬 적게 사용하여 긴 여정에 완벽합니다. 그러나 초기 킥보드에는 문제가 있었습니다. 너무 단순했다는 점입니다. 직선 도로에서는 잘 처리할 수 있었지만, 경로가 구불구불하거나 복잡한 논리 (예: 특정 캐릭터에 대한 특정 규칙을 기억하는 것) 가 필요하면 길을 잃었습니다.

이 논문은 FLASH PD-SSM이라는 새로운 발명을 소개합니다. 이는 전기 킥보드에 스마트하고 모듈식인 항법 시스템을 업그레이드하여 무거운 리무진만큼 복잡한 구불구불한 길도 잘 처리하도록 만들되, 속도와 연비 효율은 잃지 않는 것과 같습니다.

다음은 이를 단순한 개념으로 분해한 방법입니다:

1. 문제: "너무 무거운" 지도

이러한 스마트 킥보드의 이전 버전 (예: PD-SSM이라는 모델) 은 여정의 모든 단계마다 거대하고 상세한 지도를 휴대함으로써 초지능이 되려고 했습니다.

문제: 긴 여정을 위해 이 거대한 지도를 휴대하는 것은 너무 많은 메모리를 차지하여 킥보드가 멀리 가기 전에 배터리가 방전되었습니다. 실제 사용에는 너무 무거워 실용적이지 않았습니다.
결과: 다른 모델들 (예: Mamba) 은 아주 작고 단순한 지도를 휴대하기로 선택했습니다. 그들은 빠르고 가볍지만 복잡한 퍼즐을 풀거나 복잡한 규칙을 기억할 수 없었습니다.

2. 해결책: "마법 스위치"

이 논문의 저자들인 FLASH PD-SSM은 영리한 트릭을 고안해냈습니다. 모든 단계마다 거대하고 상세한 지도를 휴대하는 대신, 미리 제작된 전문 지도들의 도구상자를 만들었습니다.

작동 방식: 여정의 각 단계에서 로봇은 현재 상황을 살펴보고 도구상자에서 그 순간에 맞는 단 하나의 완벽한 지도를 선택하기 위해 스위치를 조작합니다.
비유: 운전한다고 상상해 보세요. 모퉁이를 돌 때마다 도시 전체의 새롭고 상세한 지도를 그리는 것 (이는 시간이 매우 오래 걸리고 많은 종이를 사용함) 대신, 100 개의 미리 그려진 '지역 지도' 세트를 가지고 있습니다. 다음 블록에 필요한 지도 하나만 선택하면 됩니다.
이점: 이 스위치는 매우 빠르게 조작되며 거의 공간을 차지하지 않습니다. 이로 인해 로봇은 무거운 리무진의 복잡성(어려운 논리 퍼즐을 해결할 수 있음) 을 유지하면서도 킥보드의 속도와 가벼움을 유지할 수 있습니다.

3. 그들이 증명한 것

팀은 이 새로운 "스마트 킥보드"를 세 가지 주요 방식으로 테스트했습니다:

**논리 테스트 **(FSA Emulation) 로봇에게 일련의 논리 퍼즐 (패리티 계산이나 미로 탐색 등) 을 주었습니다. FLASH PD-SSM 은 거의 모든 퍼즐을 해결했습니다 (96% 정확도). 반면 더 단순한 모델들 (예: Mamba2) 은 크게 어려움을 겪었습니다. 이는 로봇이 실제로 복잡한 규칙을 통해 "생각"할 수 있음을 증명했습니다.
**긴 기억 테스트 **(시계열) 로봇에게 17,000 단계가 넘는 매우 긴 데이터 스트림을 분석하도록 요청했습니다. FLASH PD-SSM 은 다음에 무슨 일이 일어날지 예측하는 데 가장 정확했으며, 모든 다른 경쟁자들을 능가했습니다.
언어 테스트: 이 새로운 뇌를 언어 모델 (텍스트를 작성하는 로봇) 에 넣었습니다.
- 상태 추적: 이야기에서 물체가 상자 사이를 이동하는 것을 추적하도록 요청했을 때 (예: "빨간 공이 A 상자에서 B 상자로 이동했다"), 새로운 모델은 이전 모델들보다 최종 위치를 훨씬 잘 기억했습니다.
- 작성: 인간처럼 쓰도록 훈련되었을 때, 이 모델의 하이브리드 버전 (새로운 뇌와 기존 기술을 혼합) 은 이전의 더 단순한 뇌만 사용한 모델들보다 더 잘, 더 빠르게 작성했습니다.

4. 결론

이 논문은 FLASH PD-SSM이 현재 AI 에서 가장 큰 트레이드오프인 속도 vs 지능을 해결한다고 주장합니다.

이전: 빠르지만 "멍청한"(단순한 모델) 것과 "똑똑하지만" 느리고 비싼 것 (복잡한 모델) 사이에서 선택해야 했습니다.
이제: FLASH PD-SSM 은 산을 오를 수 있는 고속 기차와 같습니다. 현재 최상위 모델 (Mamba2) 만큼 빠르게 실행되지만 훨씬 더 복잡한 작업을 처리할 수 있으며, 모든 것이 더 적은 메모리를 사용합니다.

요약하자면, 그들은 더 크고 비싼 컴퓨터를 구축할 필요 없이 AI 모델을 동시에 가볍고, 빠르고, 똑똑하게 만드는 방법을 찾았습니다.

기술 요약: FLASH PD-SSM

문제 제기

상태 공간 모델 (SSM) 은 선형 계산 확장성과 추론 중 일정한 메모리 소비를 제공하여 트랜스포머의 확장 가능한 대안을 제시합니다. 그러나 전이 행렬의 구조에 의해 결정되는 효율성과 표현력 사이의 근본적인 트레이드오프에 직면해 있습니다.

비구조화 행렬은 최대의 표현력 (임의의 유한 상태 오토마타, 즉 FSA 를 모방할 수 있음) 을 제공하지만, prohibitive 한 계산 및 메모리 비용을 초래합니다.
구조화 행렬(예: 대각선, 스케일된 단위 행렬) 은 매우 효율적이지만 제한된 표현력으로 인해, 이론적으로 비해결 가능한 변환 반군을 가진 FSA 를 모방할 수 없습니다.
구조화 희소 행렬 (PD-SSM): 최근의 진전인 PD-SSM 은 구조화된 희소 행렬을 활용하여 FSA 수준의 표현력을 달성하면서도 효율적인 병렬 계산을 가능하게 합니다. 그러나 PD-SSM 은 중요한 병목 현상에 시달립니다: 입력 의존적 희소 행렬을 생성하는 데 $O(LN^2)$ 의 메모리 비용이 발생합니다 (여기서 $L$ 은 시퀀스 길이이고 $N$ 은 은닉 상태 크기입니다). 이 오버헤드로 인해 PD-SSM 은 Mamba2 와 같은 효율적인 경쟁자보다 현저히 느리고 (40 배) 메모리 집약적이어서 대규모 배포에 적합하지 않습니다.

방법론: FLASH PD-SSM

저자들은 PD-SSM 의 생성 병목 현상을 해결하면서도 이론적 표현력 보장을 유지하는 메모리 최적화 구조화 희소 SSM 인 FLASH PD-SSM을 제안합니다.

1. 새로운 선택 메커니즘 (하드 선택)

핵심 혁신은 PD-SSM 의 소프트 선택 메커니즘을 하드 선택 메커니즘으로 대체합니다:

기존 PD-SSM: 밀집 사전 행렬의 볼록 결합을 생성하고 결과를 희소화하기 위해 열별 하드맥스를 적용합니다. 이는 매 시간 단계마다 밀집 $N \times N$ 행렬을 물리적으로 생성해야 하므로 $O(LN^2)$ 의 메모리가 필요합니다.
FLASH PD-SSM: $K$ 개의 희소 행렬로 구성된 사전 을 사전 계산합니다. 각 시간 단계에서 입력 $u_t$ 는 선형 변환을 통해 $K$ 개의 점수를 생성합니다. 모델은 가장 높은 점수를 가진 단일 행렬 (argmax) 을 선택하여 직접 사용합니다.
복잡도 감소: 선택된 행렬이 이미 희소 (원-핫 열 선택기) 이므로, 메모리 비용이 $O(LN^2)$ 에서 **$O(LN) $**으로 감소합니다. 행렬 생성에 대한 계산 비용은$ O(LN^2K) $에서$ O(LNK)$로 감소합니다.

2. 기울기 전파

선택이 이산 연산 (argmax/하드맥스) 을 포함하므로, 모델은 역전파를 위해 기울 Annealing 을 적용한 **Straight-Through Estimators (STE)**를 사용합니다:

순방향 전달에서는 하드 선택 (argmax) 을 사용합니다.
역방향 전달에서는 선택 가중치와 사전 행렬을 통한 기울기 흐름을 허용하기 위해 하드 선택을 온화한 소프트맥스로 근사화합니다.
논문의 Proposition 2 는 온도 매개변수 $\tau \to 0$ 이 될 때 정확해지는 대리 기울기 표현식을 제공합니다.

3. 하드웨어 인식 구현

저자들은 GPU 병렬성을 활용하기 위해 커스텀 3 단계 CUDA 커널(Triton 으로 래핑됨) 을 구현합니다:

Phase A: 시퀀스의 청크 내에서 독립적인 순차 재귀를 계산합니다.
Phase B: 캐리 상태를 계산하기 위해 청크 간 의존성 (접두어 스캔) 을 집계합니다.
Phase C: 청크 상태에 캐리 보정을 적용합니다.
이 설계는 메모리 트래픽을 최소화하고 대역폭 활용을 극대화하여 거의 정점에 가까운 GPU 메모리 대역폭을 달성합니다.

4. 아키텍처 통합

FLASH PD-SSM 은 멀티헤드 SSM 레이어가 있는 표준 프리-노름 아키텍처에 내장됩니다. 멀티헤드 설계는 밀집 사전 행렬에 대한 파라미터 비용을 상쇄할 수 있게 합니다. 저자들은 사전 크기 $K$ 를 $O(\sqrt{D})$ (여기서 $D$ 는 임베딩 차원) 로 스케일링하면 모델 크기를 "폭발"시키지 않으면서 파라미터 수를 관리 가능하게 유지할 수 있음을 보여줍니다.

주요 기여

알고리즘적 개선: 이산 선택 메커니즘을 통해 구조화 희소 SSM 생성의 메모리 복잡도를 $O(LN^2)$ 에서 $O(LN)$으로 감소시켜 대규모 훈련을 가능하게 합니다.
이론적 표현력: Proposition 1 을 통해 FLASH PD-SSM 이 비구조화 SSM 의 이론적 보장을 유지함을 증명합니다: 단일 $N$ 크기 레이어를 사용하여 $N$ 개의 상태를 가진 임의의 결정론적 FSA 를 정확히 표현할 수 있습니다.
하드웨어 최적화: 최첨단 처리량 및 메모리 효율성을 달성하는 커스텀 CUDA 구현을 제공합니다.
하이브리드 모델 통합: FLASH PD-SSM 이 하이브리드 LLM 에서 효과적인 플러그인 대체재로 기능하여 처리량을 희생하지 않으면서 상태 추적 능력을 향상시킨다는 것을 보여줍니다.

실험 결과

합성 상태 추적 및 표현력

FSA 모방: 합성 작업 (패리티, 모듈러 카운팅, 사이클 탐색) 에서 FLASH PD-SSM 은 **평균 정확도 96.3%**를 달성하여 대각선 SSM(Mamba2: 58.9%, Mamba3: 90.7%) 과 DPLR 모델 (Gated DeltaProduct: 80.7%) 을 크게 능가합니다. 이는 이론적 표현력의 실제 실현을 확인시켜 줍니다.
기작적 능력 (MAD 벤치마크): FLASH PD-SSM 은 **69.2%**를 기록하여 Mamba(69.3%) 와 유사한 수준으로, 강력한 컨텍스트 내 회상 및 기억 능력을 나타냅니다.

런타임 및 메모리 효율성

처리량: 시퀀스 길이가 5120 일 때, FLASH PD-SSM 은 DeltaNet 보다 18% 빠르고 Mamba2 보다 4% 빠릅니다.
메모리: 시퀀스 길이가 4000 이상일 때, DeltaNet 과 Mamba2 모두에 비해 피크 메모리 소비를 12% 이상 감소시킵니다.
확장성: 모델은 다양한 시퀀스 길이에 걸쳐 높은 처리량을 유지하며, 이전에 PD-SSM 을 방해했던 병목 현상을 해결합니다.

자연어 상태 추적

하이브리드 모델 (Granite4, Nemotron-H) 을 사용하여 Box 데이터셋(26 개 상태 전이로 확장) 에서 평가되었습니다.
이러한 모델에 삽입된 FLASH PD-SSM 레이어는 Mamba2 대비 최대 10%, DeltaNet 대비 **8%**의 분포 외 (OOD) 정확도를 향상시키면서 동등한 처리량을 유지했습니다.

언어 모델링

DELTA-PD: FLASH PD-SSM 과 DeltaNet 을 결합한 하이브리드 모델 (1:3 비율) 이 FineWeb100B 에서 사전 훈련되었습니다.
성능: DELTA-PD 는 순수 DeltaNet 을 정확도 (평균 42.3% 대 41.3%) 와 처리량(253k 토크/초 대 168k 토크/초, 1.5 배 개선) 에서 모두 능가했습니다.
순수 FLASH PD-SSM 보다도 우수하여, 하이브리드화가 서로 다른 SSM 유형의 강점을 균형 있게 조화시킨다는 것을 시사합니다.

시계열 모델링

시퀀스 길이가 17,000 이상인 UEA 다변량 시계열 분류 작업에서 최첨단 (SoTA) 정확도를 달성했습니다.
PPG-DaLiA 예측 작업에서 거의 최첨단 평균 제곱 오차 (MSE) 를 달성했습니다.

중요성 및 주장

이 논문은 FLASH PD-SSM 이 SSM 을 제한해 온 효율성 - 표현력 트레이드오프를 성공적으로 깨뜨렸다고 주장합니다. 구조화 희소 행렬의 메모리 오버헤드를 줄임으로써, 현대 언어 모델에 필요한 규모에서 FSA 모방이 가능한 매우 표현력 있는 모델의 사용을 가능하게 합니다.

저자들은 FLASH PD-SSM 을 단순히 이론적 개선이 아닌, 최전선 하이브리드 모델의 기존 SSM 레이어를 위한 실용적인 플러그인 대체재로 위치시킵니다. 이는 성능 지형에서 독특한 위치를 차지합니다: Figure 1 의 "우상단 사분면"에서 식별된 유일한 SSM 으로, 동시에 매우 표현력 있고 빠르며 메모리 효율적입니다.

논문은 Straight-Through Estimators 의 사용과 관련된 한계를 인정하며, 고도로 이산적인 순차 방법에서의 그 효과에 대한 이론적 설명은 향후 연구의 열린 영역으로 남아 있음을 지적합니다. 그러나 합성, 기작적, 자연어 작업 전반에 걸친 경험적 결과는 이 접근법을 확장 가능하고 표현력 있는 시퀀스 모델링을 위한 강력한 솔루션으로 검증합니다.