[Paper] MAR: 모듈 인식 아키텍처 정제를 통한 효율적인 대규모 언어 모델

Source: arXiv - 2601.21503v1

개요

대형 언어 모델(LLM)은 많은 AI 제품의 핵심이 되었지만, 이들의 2차 시간 복잡도를 갖는 어텐션과 조밀한 피드‑포워드 네트워크(FFN)는 추론 시 계산량과 에너지 소모가 크게 발생합니다. 논문 **“MAR: Efficient Large Language Models via Module‑aware Architecture Refinement”**는 LLM의 비용이 많이 드는 부분을 보다 효율적인 대안—시퀀스 처리를 위한 상태 공간 모델(SSM)과 희소화된 활성화 함수 기반 FFN—으로 교체하는 두 단계 프레임워크를 제시하면서 성능을 유지하거나 오히려 향상시킵니다.

핵심 기여

• Module‑aware Architecture Refinement (MAR): 2차 곱셈 주의(attention)를 선형 시간 SSM으로 교체하고, 각 레이어를 손수 튜닝하지 않고도 FFN 활성화를 희소화하는 체계적인 파이프라인.
• Adaptive Ternary Multi‑step Neuron (ATMN): SSM과 Spiking Neural Networks (SNN) 사이의 시간 불일치를 연결해, 저정보 밀도 신호를 효율적으로 처리할 수 있게 하는 새로운 스파이킹 뉴런 설계.
• Spike‑aware Bidirectional Distillation Strategy (SBDS): 조밀한 교사 모델의 지식을 SSM과 SNN 모듈 모두에 동시에 증류하는 학습 레시피로, 정제된 아키텍처가 원래 정확도를 회복하도록 보장.
• Comprehensive energy‑aware evaluation: 저자들은 실제 하드웨어에서 추론 에너지를 측정하여, 조밀한 베이스라인 대비 **45 %**까지 에너지 감소를 달성하면서 BLEU/GLUE 점수를 동일하거나 상회함을 보여줌.
• Scalable to large model sizes: MAR는 파라미터 수가 비슷하거나 더 큰 다른 “효율적인” LLM 변형(예: LoRA‑pruned, quantized, 혹은 sparsified 모델)보다도 뛰어난 성능을 발휘.

방법론

2단계 정제

1. 1단계 – 어텐션 교체: 각 Transformer 자체 어텐션 블록을 선형 시간 SSM(예: HiPPO 기반)으로 교체합니다. SSM은 토큰 시퀀스를 O(N) 시간에 처리하여 N² 어텐션 행렬을 없앱니다.
2. 2단계 – FFN 희소화: 활성화 기반 프루닝을 사용해 밀집 FFN을 희소화된 버전으로 변환합니다. 토큰당 상위 k 활성화만 유지하여 밀집 행렬 곱셈을 훨씬 저렴한 희소 연산으로 바꿉니다.

스파이킹 인식 통합

• ATMN은 연속적인 SSM 출력을 여러 타임스텝에 걸쳐 삼진 스파이크(‑1, 0, +1)로 변환하여 정보를 보존하면서 다운스트림 SNN이 최소 에너지로 동작하도록 합니다.
• SBDS는 양방향 지식 증류를 수행합니다: 밀집 교사가 SSM의 은닉 상태를 안내하고, SNN의 스파이킹 동역학은 교사의 토큰 수준 로짓을 모방하도록 정규화됩니다. 이 공동 학습은 구조 교체로 인한 성능 손실을 회복합니다.

학습 파이프라인

• 표준 밀집 LLM(교사)을 사전 학습합니다.
• 교사를 고정하고 MAR‑정제 학생을 SBDS와 함께 학습하며, SSM 전용과 SSM+SNN 순전파를 교대로 수행합니다.
• 수천 스텝마다 업데이트되는 크기 기반 마스크를 사용해 희소화된 FFN을 미세 조정합니다.

전체 과정은 레이어 깊이와 토큰별 정보 밀도를 기준으로 SSM을 적용할지 전통 어텐션을 적용할지를 결정하는 “모듈 인식” 스케줄러를 통해 자동화됩니다.

결과 및 발견

• 성능 동등성: MAR는 언어 이해 벤치마크(GLUE, SuperGLUE)와 생성 작업(BLEU, ROUGE) 전반에 걸쳐 dense 모델 정확도의 >99 %를 회복합니다.
• 에너지 절감: NVIDIA A100 및 저전력 ARM CPU에서의 실제 측정 결과, 토큰당 추론 에너지가 최대 45 % 감소함을 보여줍니다.
• 확장성: 13 B 파라미터로 확장할 때에도 MAR는 더 큰 dense 베이스라인보다 여전히 우수하며, 효율성 향상이 모델 크기와 함께 누적됨을 나타냅니다.

Practical Implications

• Edge‑AI & on‑device LLMs: 선형‑시간 SSM 및 스파이킹 구성 요소 덕분에 배터리 제한이 있는 장치(예: 스마트폰, 웨어러블)에서도 품질을 희생하지 않고 정교한 언어 모델을 실행할 수 있습니다.
• Cloud cost reduction: 추론이 많이 필요한 서비스(챗봇, 코드 어시스턴트)에서 MAR은 GPU 사용 시간 비용과 탄소 발자국을 줄여 운영 비용을 직접 낮출 수 있습니다.
• Simplified deployment pipelines: MAR이 attention/FFN 모듈을 바로 대체할 수 있기 때문에 기존 Transformer 코드베이스를 최소한의 엔지니어링 작업으로 재구성할 수 있습니다—모듈을 교체하고 제공된 SBDS 학습 스크립트를 실행하면 됩니다.
• Compatibility with other efficiency tricks: MAR은 양자화, 모델 병렬화, 또는 LoRA 파인튜닝과 결합될 수 있어 최적화를 위한 계층적 접근 방식을 제공합니다.

제한 사항 및 향후 작업

• SSM 및 스파이크에 대한 하드웨어 지원: 저자들은 GPU/CPU에서 에너지를 측정했지만, 가장 큰 이득은 삼진 스파이크를 네이티브로 처리하는 특수 신경형 또는 ASIC 가속기에서 나타납니다; 보다 넓은 하드웨어 지원은 아직 초기 단계입니다.
• 학습 오버헤드: 2단계 정제와 양방향 증류를 결합하면 일반적인 밀집 모델에 비해 약 30 %의 추가 학습 시간이 소요됩니다.
• 시간 정렬 민감도: ATMN 설계는 특정 시퀀스 길이를 가정합니다; 매우 긴 문서(> 4 k 토큰)는 여전히 잔여 시간 불일치로 인해 영향을 받을 수 있습니다.

향후 방향은 다음과 같습니다:

1. 신흥 AI 칩을 위한 SSM 커널 공동 설계.
2. 런타임 지연 제약에 반응하는 적응형 희소성 스케줄.
3. 멀티모달 트랜스포머(비전‑언어, 음성)로 MAR 확장.

핵심 요약: MAR은 개발자가 의존하는 사용자 경험 품질을 손상시키지 않으면서 대형 언어 모델을 더 친환경적이고 빠르게 만들 수 있는 실용적인 경로를 제공합니다. 선형 시간 역학과 스파이킹 희소성 관점에서 핵심 어텐션 및 피드‑포워드 블록을 재고함으로써, 이 프레임워크는 클라우드 클러스터에서 엣지까지 진정으로 확장 가능한 LLM 배포의 문을 엽니다.

저자

• Junhong Cai
• Guiqin Wang
• Kejie Zhao
• Jianxiong Tang
• Xiang Wang
• Luziwei Leng
• Ran Cheng
• Yuxin Ma
• Qinghai Guo

논문 정보

• arXiv ID: 2601.21503v1
• 분류: cs.AI, cs.CL, cs.LG, cs.NE
• 출판일: 2026년 1월 29일
• PDF: PDF 다운로드