NAS-Driven Hardware Accelerator Exploration for Edge AI and Quantization Effects on the Pareto Space

당신은 궁극의 로봇 두뇌를 만들려고 노력 중이라고 상상해 보세요. 하지만 엄격한 규칙이 하나 있습니다. 반드시 스마트워치나 드론 같은 아주 작은, 배터리로 작동하는 기기 안에 들어가야 한다는 것입니다. 이것이 바로 '에지 AI(Edge AI)'의 세계입니다. 거대하고 전력을 많이 소모하는 클라우드 서버 대신 작은 기기 속에 인공지능이 살아가는 곳이죠. 이를 실현하기 위해 엔지니어들은 까다로운 퍼즐을 풀어야 합니다. 얼굴이나 소리를 인식할 수 있을 만큼 똑똑하면서도, 배터리를 몇 분 만에 다 써버리지 않을 만큼 작고 빨라야 하는 두뇌가 필요합니다.

완벽한 두뇌 설계를 찾기 위해 과학자들은 '신경망 구조 탐색(Neural Architecture Search, NAS)'이라는 방법을 사용합니다. 이것은 마치 수백만 개의 서로 다른 레시피(신경망 설계)를 시도하여, 가장 맛있는 맛(최고의 정확도)을 내면서도 재료(연산 능력)를 가장 적게 사용하는 레시피를 찾아내는 초고속 자동 셰프와 같습니다. 보통 이 셰프들은 '전정밀도(full-precision)' 재료로 요리합니다. 이는 마치 정확하고 무거운 계량컵을 사용하는 것과 같습니다. 하지만 아주 작은 기기를 위해서는 '양자화(quantization)'로 전환해야 합니다. 이는 더 작은, 가벼운 계량컵(숫자를 표현하는 데 더 적은 비트를 사용)으로 바꾸는 것과 같습니다. 여기서 큰 질문이 생깁니다. 만약 계량컵을 바꾼다면, 주방에서 완벽해 보였던 레시피가 길 위에서도 똑같은 맛을 낼 수 있을까요? 이 논문은 바로 그 질문을 파고듭니다. 측정 시스템을 바꾸는 것이 최적의 로봇 두뇌를 찾는 과정에 어떤 영향을 미치는지, 그리고 그 두뇌를 특수한 재구성 가능 칩에 어떻게 맞출 수 있는지 탐구합니다.

---

작은 두뇌를 향한 3단계 여정

이 논문의 저자들은 AI를 하드웨어에 맞추는 문제를 해결하기 위해 영리한 3단계 파이프라인을 제안합니다. 그들은 이를 "NAS-then-quantized(NAS 후 양자화)" 접근 방식이라고 부르는데, 이는 컴퓨터로 자동차를 먼저 설계한 다음, 울퉁불퉁한 흙길에서 어떻게 주행하는지 테스트하고, 마지막으로 지형에 맞는 엔진을 만드는 것과 비슷합니다.

1단계: 하드웨어 불가지론적 프런트엔드 (The Hardware-Agnostic Frontend)
먼저, 팀은 실제 모델을 구축하지 않고도 설계의 우수성을 예측하는 똑똑한 예측 모델인 '대리 모델(surrogate model)'을 사용합니다. 그들은 NAS-Bench-201이라는 데이터셋에 포함된 15,625개의 다양한 신경망 설계 라이브러리를 바탕으로 이 예측 모델을 학습시킵니다. 이 단계는 "하드웨어 불가지론적"입니다. 즉, 아직 특정 칩을 고려하지 않으며, 오직 정확도와 필요한 수학 연산량(FLOPs)을 기준으로 최적의 설계를 찾습니다. 그들은 두 가지 탐색 전략을 사용합니다: 단순한 '무작위 탐색(Random Search)'과 자연 선택을 모방하여 더 나은 설계를 진화시키는 '다목적 진화 알고리즘(Multi-Objective Evolutionary Algorithm)'입니다. 이 단계를 거치면 10개의 상위 후보군 리스트가 만들어집니다.

2단계: 양자화 브릿지 (The Quantization Bridge)
다음은 현실 점검 단계입니다. 팀은 이 10개의 상위 후보를 가져와 '사후 학습 양자화(Post-Training Quantization, PTQ)' 기술을 사용하여 크기를 줄입니다. 구체적으로는 일반적인 32비트 대신 4비트(INT4)만을 사용하도록 전환합니다. 이는 고화도 사진을 오래된 휴대폰에 맞게 압축하는 것과 같습니다. 그들은 이를 위해 Brevitas라는 도구를 사용합니다.

여기서 흥미로운 점이 발생합니다. 모델을 압축하면, 첫 번째 단계에서 훌륭해 보였던 설계들이 실패하거나 성능이 떨어질 수 있습니다. '양자화 브릿지'는 일종의 필터 역할을 합니다. 압축을 견뎌내는 정도에 따라 설계의 순위를 다시 매깁니다. 만약 설계가 무너지면 탈락합니다. 만약 살아남은 설계가 하나도 없다면, 시스템은 다시 1단계로 돌아가 처음부터 다시 시도합니다. 그들이 INT4를 선택한 이유는, 매우 공격적인 방식임에도 불구하고 약간의 추가 학습(단 2 에포크)만 있으면 모델이 정확도를 크게 회복할 수 있다는 것을 발견했기 때문입니다. 정확도가 최저 9.49%에서 83.53%까지 뛰어오른 것은 원래 모델의 정보가 상당 부분 보존되었음을 증명합니다.

3단계: 하드웨어 인지 백엔드 (The Hardware-Aware Backend)
마지막으로, 살아남은 설계들은 집이 필요합니다. 팀은 이 압축된 두뇌들을 CGRA(Coarse-Grained Reconfigurable Array)라고 불리는 특정 유형의 하드웨어에 매핑합니다. CGRA는 어떤 모양의 블록이든 맞출 수 있는 레고 판과 같습니다. 그들은 하드웨어 배치(프로세싱 유닛 수, 메모리 깊이, 데이터 경로 변경)를 탐색하기 위해 진화 알고리즘을 사용합니다. 물리적인 칩을 직접 만드는 대신, 수학적 시뮬레이터인 '자유 분석 오라클(free analytical oracle)'을 사용하여 각 배치가 얼마나 빠르고 효율적인지 예측합니다. 그들은 지연 시간을 최소화하고 하드웨어 자원 활용을 극대화하는 구성을 선택합니다.

큰 놀라움: 옛날 지도를 버리지 마세요

이 논문에서 가장 흥밀한 발견은 단순히 이 3단계 파이프라인 자체가 아니라, '전정밀도(full-precision)' 세계와 '압축된' 세계 사이의 관계에 대해 알아낸 점입니다.

보통 모델을 압축할 때, 압축된(INT4) 세계에 특화된 새로운 예측 모델을 처음부터 다시 학습시켜야 한다고 생각할 수 있습니다. 저자들은 이 아이디어를 테스트했습니다. 그들은 두 가지 접근 방식을 비교했습니다:

1. 전용 INT4 대리 모델: 압축된 데이터에 특화되어 학습된 예측 모델.
2. FP32 제로샷 대리 모델: 원래의 전정밀도 데이터로 학습된 후, 재학습 없이 압축된 모델의 성능을 예측하는 데 사용된 예측 모델.

결과는 놀라웠습니다. 시뮬레이션 결과, FP32 제로샷 대리 모델이 전용 INT4 모델보다 더 우수한 성능을 보였습니다. 이 모델을 사용하여 탐색 전략을 실행했을 때, FP32 예측 모델은 INT4 전용 예측 모델보다 더 넓은 범위의 최적의 트레이드오프(파레토 공간)를 커버하는 설계들을 찾아냈습니다.

왜 그럴까요? 저자들은 전정밀도 데이터가 훨씬 더 "깨끗하고" 노이즈가 적기 때문이라고 설명합니다. 비록 압축된 세계가 다르긴 하지만, 어떤 설계가 "좋고" 어떤 설계가 "나쁜지"에 대한 순위는 놀라울 정도로 유사하게 유지됩니다. 논문은 비록 설계의 정확한 순위는 변하지만(약 21%의 경우, 전정밀도에서 더 좋았던 설계가 압축 후에 더 나빠짐), '최고'의 설계들이 가진 전체적인 구조는 충분히 안정적이어서 기존의 지도(FP32 예측 모델)가 여전히 신뢰할 수 있는 가이드가 된다고 언급합니다.

퍼즐의 안정성

팀은 단순히 추측한 것이 아니라, 혼돈의 정도를 측정했습니다. 그들은 15,625개의 모든 아키텍처를 조사하여 INT4로 전환했을 때 '파레토 프런트(Pareto front, 최적의 설계 목록)'가 얼마나 이동하는지 확인했습니다.

• 변화(The Shift): 상위 설계 목록이 완전히 재편되었습니다. 원래의 상위 설계 중 압축된 세계에서 상위로 살아남은 것은 하나도 없었습니다.
• 역전(The Flip): 약 21.73%의 경우, 두 설계 간의 관계가 뒤집혔습니다(전정밀도에서는 설계 A가 B보다 좋았으나, 압축 후에는 B가 A보다 좋아짐).
• 상관관계(The Correlation): 이러한 혼돈에도 불구하고, 전체적인 순위 상관관계는 0.6655로 여전히 중간 수준을 유지했습니다. 즉, 좋은 설계들의 전반적인 "형태"는 보존되었습니다.

또한 그들은 더 작고 단순한 모델(수학 연산이 적은 모델)이 압축에 훨씬 더 민데이션(민감)하며, 종종 정확도를 잃는다는 것을 발견했습니다. 반면, 더 크고 복잡한 모델은 구조적으로 작은 숫자가 유발하는 오류를 자연스럽게 완화해주기 때문에 놀라울 정도로 견고했습니다.

하드웨어 결과

마지막으로, 그들은 가장 잘 살아남은 세 가지 설계를 CGRA 하드웨어에 매핑했습니다. 그들은 설계가 서로 다름에도 불구하고, 모두 정확히 동일한 최적의 하드웨어 구성으로 수렴한다는 것을 발견했습니다. 이는 16행 66열의 프로세싱 유닛 그리드와 특정 메모리 깊이를 가진 구성입니다. 이는 특정 신경망 설계가 유사한 빌딩 블록을 사용한다면, 해당 작업에 대한 하드웨어의 '스윗 스팟(sweet spot)'은 매우 일관적이라는 것을 시사합니다.

결론

이 논문은 작은 기기를 위한 AI 모델을 축소할 때마다 매번 바퀴를 새로 발명할 필요가 없다는 것을 보여줍니다. 전정밀도에서 최적의 설계를 찾고, 압축 테스트를 통해 필터링한 다음, 유연한 하드웨어에 매핑하는 3단계 과정을 통해 정확도와 크기 사이의 완벽한 균고를 효율적으로 찾을 수 있습니다. 가장 중요한 것은, 전정밀도 검색 도구가 압축된 세계에서도 우리를 안내할 수 있다는 믿음을 줌으로써, 강력한 AI를 우리 주머니 속으로 넣으려는 여정의 시간과 노력을 절약할 수 있다는 점입니다.