Fast and Memory-Efficient Wavelet Convolutions via I/O-Aware Reformulation

당신이 사진을 보고 그 안에 무엇이 들어있는지 정확하게 말할 수 있는 초지능 로봇을 만들려고 한다고 상상해 보세요. 이를 위해 로봇은 단순히 작은 점 하나가 아니라 전체 그림을 한 번에 '볼' 수 있어야 합니다. 컴퓨터 과학의 세계에서는 이를 '수용 영역(receptive field)'이 크다고 부릅니다. 오랫동안 이 넓은 시야를 로봇에게 제공하는 가장 좋은 방법은 작은 필터들을 여러 층으로 쌓아 올리는 것이었습니다. 마치 렌즈를 높게 쌓아 올려 탑을 만드는 것과 같습니다. 하지만 이 탑은 매우 빠르게 무거워지고 느려집니다.

최근 과학자들은 '웨이브릿 컨볼루션(Wavelet Convolutions, 또는 WTConv)'이라는 영리한 지름길을 발견했습니다. 이 방법은 렌즈를 쌓는 대신, '웨이브릿 변환'이라는 수학적 마법 기술을 사용하여 규칙의 수를 아주 작게 유지하면서도 전체적인 큰 그림을 볼 수 있게 해줍니다. 이는 마치 단 하나의 창문만으로도 도시 전체를 볼 수 있는 망원경을 가진 것과 같습니다. 문제는 무엇일까요? 이 방법은 수학적으로는 매우 훌ert하지만, 이를 실행하는 컴퓨터가 데이터를 너무 많이 옮기고 있었다는 점입니다. 그것은 마치 사서가 책 한 권을 가져올 때마다 매번 지하로 내려갔다 올라오는 것과 같았습니다. 책을 바로 옆 선반에서 꺼내는 대신 말이죠. 이로 인해 로봇은 믿을 수 없을 정도로 느려졌고, 메모리를 엄청나게 잡아먹으며 그 잠재력을 낭비했습니다.

"I/O-Aware Reformulation을 통한 빠르고 메모리 효율적인 웨이브릿 컨볼루션"이라는 제목의 이 논문은 바로 그 문제를 다룹니다. 벤구리온 대학교의 연구진인 저자들은 속도 문제가 수학이 어려워서가 아니라, 컴퓨터가 메인 메모스에서 데이터를 너무 많이 옮기고 다니는 데서 발생했다는 것을 깨달았습니다. 그들은 이 '웨이브릿 트릭'의 새로운 초효율 버전을 구축하여, 데이터가 프로세서 자체의 칩 위에 머물도록 했습니다. 이렇게 함으로써 그들은 단순히 로봇을 조금 더 빠르게 만든 것이 아니라, 느릿한 과정을 전력 질주로 바꾸어 놓았습니다. 그들의 새로운 방식은 기존 버전보다 최대 4.35배 더 빠르며 메모리는 절반 미만으로 사용합니다. 가장 인상적인 점은, 이들이 대체하려 했던 표준 방식보다도 더 뛰어난 성능을 보여주었다는 것입니다. 이는 데이터의 스마트한 재구성이 새로운 발명만큼이나 강력할 수 있음을 증명합니다.

문제점: "지하로 달려가는" 사서

저자들이 무엇을 했는지 이해하기 위해, 책(데이터)은 거대한 지하(고대역폭 메모리, HBM)에 저장되어 있고, 독서 테이블(프로세서)은 꼭대기 층에 있는 도서관을 상상해 보세요. 기존의 웨이브릿 컨볼루션 방식은 사서가 계산을 할 때마다 지하로 달려가 책을 가져와서, 간단한 수학 문제를 풀고, 다시 책을 갖다 놓은 뒤, 다음 문제를 위해 또 다시 내려가는 것과 같았습니다.

수학 문제 자체는 간단했지만, 사서는 계산을 하는 시간보다 데이터를 옮기는 데 90%의 시간을 허비했습니다. 저자들은 기존 방식이 데이터 한 조각당 메모리 시스템을 통해 약 18회에서 21회 정도 이동시킨다는 것을 계산해 냈습니다. 너무 비효식적이어서 컴퓨터는 실제로 생각하기보다는 데이터가 도착하기를 기다리는 '메모리 바운드(memory-bound)' 상태였습니다. 그들은 컴퓨터가 이 교통 체증에 갇혀 있어서 잠재 속도의 약 3%만을 사용하고 있다는 것을 발견했습니다.

해결책: 세 가지 마법의 기술

저자들은 새로운 수학을 발명한 것이 아니라, 단지 수학을 수행하는 방법을 바꿨습니다. 그들은 사서가 지하로 달려가는 것을 멈추기 위해 세 가지 특정 기술을 사용했습니다.

1. "즉석 처리" 기술 (Recomputing Analysis)
기존 방식에서는 컴퓨터가 먼저 데이터를 특수한 형식(이를 "Haar 분석"이라 부릅니다)으로 변환한 다음, 그 결과를 지하에 저장하고 다시 돌아와 사용했습니다. 저자들은 이 변환 과정이 매우 가볍다는 것을 깨달았습니다. 그것은 그저 숫자를 더하고 빼는 작업일 뿐이었습니다. 그래서 그들은 결과를 저장하지 않기로 했습니다. 대신 컴퓨터에게 이렇게 명령했습니다. "이것을 적어두지 마세요. 그냥 지금 당장 프로세서 안에서 직접 계산하세요." 이것은 사서가 메모장에 적어서 지하에 저장하러 가는 대신, 머릿속으로 직접 계산하기로 결정한 것과 같습니다. 이 덕분에 엄청난 양의 왕복 이동을 줄일 수 있었습니다.

2. "원패스" 기술 (Collapsing the Synthesis)
기존 방식은 최종 이미지를 단계별로 구축했습니다. 조각 하나를 가져와서 다음 조각과 더하고, 결과를 저장하고, 그 결과에 다시 다음 조각을 더하고 저장하는 식이었습니다. 이는 벽돌 하나를 놓고 지하로 내려가 다음 벽돌을 가져와 다시 놓는 과정을 반복하며 탑을 쌓는 것과 같았습니다. 저자들은 단 한 번의 과정으로 최종 결과를 계산할 수 있는 수학 공식을 찾아냈습니다. 지하를 오가는 대신, 주소에 기반하여 모든 벽돌이 어디로 가야 하는지 파지도를 보고 한 번에 배치하는 것입니다. 이를 통해 중간 단계의 "탑"을 저장하고 다시 불러올 필요를 없앴습니다.

3. "사전 혼합" 기술 (Folding Scales)
마지막으로, 기존 방식은 데이터에 "스케일(곱하기 값)"을 적용하는 별도의 단계를 거쳤는데, 이는 데이터를 읽고, 곱하고, 다시 쓰는 데 또 다른 지하 이동을 의미했습니다. 저자들은 숫자를 곱하는 것이 필터 자체의 숫자를 바꾸는 것과 같다는 것을 깨달았습니다. 그래서 프로세스가 시작되기도 전에 스케일을 필터 가중치에 미리 섞어버렸습니다. 이는 커피 가루에 설탕을 미리 섞어두어, 나중에 따로 설탕을 넣기 위해 멈출 필요가 없게 만든 것과 같습니다. 이로써 프로세스에서 단계 하나를 완전히 제거했습니다.

결과: 달팽이 대신 로켓선

저자들이 이 세 가지 기술을 결차했을 때, 결과는 극적이었습니다. 그들은 강력한 컴퓨터 칩(RTX A6000)에서 새로운 "퓨즈드(Fused)" 버전을 기존의 "레퍼런스(Reference)" 버전과 비교 테스트했습니다.

• 속도: 가장 까다로운 시나리오(신경망 학습)에서, 새로운 버전은 표준 정밀도(fp32)에서 기존 버전보다 3.71~4.35배 더 빨랐으며, 반 정밀도(fp16)에서는 2.68~3.09배 더 빨랐습니다.
• 메모리: 메모리 사용량을 약 1.83~2.31배 줄였습니다. 이는 컴퓨터가 공간 부족 없이 더 큰 이미지나 더 복잡한 모델을 처리할 수 있음을 의미합니다.
• 결정적인 승리: 가장 놀라운 발견은 새로운 웨이브릿 방식이 단순히 기존의 문제를 해결한 것에 그치지 않고, 실제로 그것이 대체하려 했던 표준 방식보다 더 빨라졌다는 점입니다. 기존의 웨이브릿 방식은 일반적인 "뎁스와이즈 컨볼루션(depthwise convolution, AI의 흔한 구성 요소)"보다 느렸습니다. 하지만 저자들의 새로운 기술을 통해, 웨이브릿 방식은 학습 시 해당 표준 방식보다 1.27~1.50배 더 빨라졌습니다.

그들은 또한 새로운 방식이 정답을 바꾸지 않았는지도 확인했습니다. 수학적 순서만 바뀌었을 뿐 계산 방식은 동일했기에, 로봇은 여전히 올바르게 학습했습니다. 그들은 다양한 이미지 크기, 다양한 레이어 수, 그리고 다른 유형의 컴퓨터 칩(NVIDIA RTX PRO 6000)에서도 테스트를 진행했으며, 속도 향상은 모든 곳에서 일관되게 나타났습니다.

이것이 왜 중요한가

이 논문은 우리에게 귀중한 교훈을 줍니다. 어떤 수학적 아이디어가 이론적으로는 효율적이라 할지라도, 그것이 실제 세상에서 빠르게 작동한다는 보장은 없다는 점입니다. 만약 컴퓨터가 생각하는 대신 데이터를 옮기는 데 시간을 허비하고 있다면, 세상에서 가장 훌륭한 수학도 도움이 되지 않습니다. 데이터가 이동하는 "배관 구조"를 살펴보고, 데이터를 프로세서 근처에 머물게 하도록 프로세스를 재설계함으로써, 저자들은 느리고 메모리를 많이 잡아먹는 도구를 번개처럼 빠른 도구로 탈바꿈시켰습니다. 그들은 복잡한 다단계 프로세스에서 때로는 더 빠른 엔진을 만드는 것보다, 자동차가 교통 체증에 갇히지 않게 만드는 것이 더 효과적이라는 것을 보여주었습니다.