OpenAI 블록-스파스 GPU 커널

OpenAI는 블록-스파스 가중치를 가진 신경망을 위해 설계된 고도로 최적화된 GPU 커널을 출시했습니다. 이러한 커널은 동일한 파라미터 및 계산 예산 내에서 전통적인 밀집 네트워크보다 훨씬 넓고 깊은 모델의 훈련을 가능하게 하며, 종종 cuBLAS 또는 cuSPARSE보다 주문 단위 이상 빠른 속도로 실행됩니다.

블록-스파스 커널의 기술적 구현

블록-스파스 커널은 개별 가중치가 아닌 블록 수준에서의 희소성을 고려하여 선형 연산을 최적화합니다. 이 접근 방식은 어떤 블록이 zero일 경우 계산을 건너뛰므로 계산 및 저장 비용을 줄입니다.

주요 사양

지원되는 레이어: 커널은 완전 연결 및 컨볼루션 레이어에서 사용되도록 설계되었습니다. 컨볼루션 레이어에서는 입력 및 출력 피처 차원에 희소성이 적용되며, 공간적 연결성은 영향을 받지 않습니다.
최적화된 블록 크기: 구현은 8x8, 16x16, 32x32 블록 크기에 최적화되었습니다.
구성 가능성: 블록 수준에서의 희소성 패턴은 완전히 구성 가능합니다.
효율성: 계산 비용과 파라미터 저장 비용은 입력/출력 피처의 총 수가 아닌 0이 아닌 가중치의 수에만 비례합니다.

성능 및 벤치마크

NVIDIA Titan X Pascal GPU에서 CUDA 8을 사용하여 수행된 테스트는 표준 라이브러리보다 상당한 속도 향상을 보여주었습니다. 12,288개의 히든 유닛을 가진 넓은 상태, 블록 크기 32x32, 미니배치 크기 32에서 커널은 cuBLAS보다 성능이 뛰어났으며, cuSPARSE와 비교했을 때 더욱 큰 속도 향상이 관찰되었습니다.

응용 및 연구 결과

OpenAI는 이러한 커널을 활용하여 여러 아키텍처 패턴을 탐구했으며, 특히 생물학적 뇌의 연결성을 모방하여 높은 희소성 속에서도 정보를 빠르게 전파할 수 있게 하는 'small-world' 신경망을 주로 연구했습니다.

Small-World LSTMs

small-world 희소 연결성을 구현함으로써 OpenAI는 약 20,000개의 히든 유닛을 가진 LSTM을 훈련시켰는데, 이는 유사한 파라미터 수를 가진 밀집 네트워크보다 다섯 배 더 넓습니다. 이 아키텍처는 텍스트 생성 모델링 및 반감독 감정 분류 결과를 향상시켰습니다.

감정 표현 학습

동일한 파라미터 수를 가진 밀집 가중치 행렬과 블록-스파스 변형을 비교한 실험에서, 희소 모델은 모든 감정 데이터셋에서 밀집 모델보다 성능이 뛰어났습니다. 구체적으로, 문서 수준 IMDB 데이터셋에서 희소 모델은 오류율을 5.91% (Miyato et al, 2016)에서 5.01%로 감소시켰습니다.

이미지 및 텍스트 압축

텍스트 모델링: 약 1억 개의 파라미터를 가진 모델에서 문자당 비트 수가 1.059에서 1.048로 감소했습니다.
이미지 모델링: CIFAR-10용 수정된 PixelCNN++ 모델에서 일반 2D 컨볼루션 커널을 희소 커널로 대체함으로써 OpenAI는 다른 하이퍼파라미터를 고정한 채 네트워크를 깊게 하여 차원당 비트 수를 2.92에서 2.90로 감소시켰습니다.

미래 연구 방향

OpenAI는 블록-스파스 커널을 사용하여 추가 탐구의 세 가지 주요 영역을 제시합니다:

추론 속도 향상: 이러한 커널과 사후 훈련 가지치기를 결합하여 추론 시 가능한 벽시계 시간 속도 향상을 조사합니다.
동적 희소성 구조: 생물학적 뇌의 시냅스 가지치기와 유사하게 개발 중에 그래디언트를 사용하여 최적의 희소성 구조를 학습할 수 있는지 탐구합니다.
확장 응용: 더 작은 밀집 모델과 동일한 파라미터 및 계산 비용을 유지하면서巨大한 가중치 행렬을 가진 모델을 훈련할 수 있는 능력이 가장 큰 성능 향상을 제공하는 특정 응용 분야를 식별합니다.