Capsule Networks (CapsNet)

캡슐 네트워크 (CapsNets)는 캡슐로 구성된 신경망입니다. 캡슐은 출력이 단일 숫자가 아니라 벡터인 작은 뉴런 그룹입니다. 캡슐 벡터의 길이는 숫자, 눈, 바퀴와 같은 개체가 존재할 확률을 나타내며, 벡터의 방향은 해당 개체의 위치, 크기, 방향, 획 두께를 비롯한 자세 및 변형 매개변수를 인코딩합니다. Geoffrey Hinton, Alex Krizhevsky, Sida Wang은 2011년 자동 인코더 변환에 관한 논문에서 캡슐을 소개했으며, Sara Sabour, Nicholas Frosst, Hinton은 캡슐 간 동적 라우팅에서 캡슐을 실용적인 분류기로 발전시켰습니다. 이 논문은 NIPS 2017에서 발표되었으며, 합의 기반 라우팅이라는 반복 알고리즘을 사용해 연속된 레이어의 캡슐을 연결합니다.

캡슐이 맥스 풀링을 대체하는 이유#

합성곱 신경망 (CNN)은 스칼라 활성화 값으로 특징을 감지하고, 맥스 풀링으로 특징 맵을 축소하는 경우가 많습니다. 맥스 풀링은 각 윈도우에서 가장 강한 응답을 유지하고 해당 응답이 윈도우 안에서 발생한 위치는 버립니다. 이 방식은 작은 이동에 대한 탐지의 강건성을 높이지만, 일부가 존재하는지만 확인하는 네트워크는 입이 눈 위에 있는 얼굴도 받아들일 수 있습니다. 이를 피카소 문제라고 합니다. Hinton은 불변성보다 등변성이 중요하다고 주장했습니다. 물체가 이동하거나 회전하면 캡슐 벡터의 길이, 즉 물체가 존재할 확률은 그대로 유지하면서 벡터의 방향은 물체와 함께 회전해야 합니다. 이렇게 시점 정보를 학습하는 캡슐은 CNN이 새로운 각도에서 물체를 인식하도록 사용하는 데이터 증강을 덜 필요로 합니다.

Ultralytics YOLO26은 캡슐 없이도 물체 위치를 유지합니다. 스트라이드 2인 합성곱 5개가 입력을 원래 해상도의 1/32로 축소하고, 유일한 맥스 풀링 블록인 SPPF는 스트라이드 1에서 5 x 5 윈도우로 세 차례 연속 풀링합니다. 이를 통해 탐지 헤드가 예측하는 그리드를 축소하지 않고 수용 영역을 넓힙니다.

동적 라우팅의 작동 방식#

하위 레이어의 각 캡슐은 출력에 학습된 가중치 행렬을 곱해 상위 레이어에 있는 모든 캡슐의 출력을 예측합니다. 라우팅은 각 하위 캡슐의 출력이 각 상위 캡슐에 얼마나 전달되는지를 결정합니다. 0에서 시작하는 라우팅 로짓에 소프트맥스를 적용해 결합 계수를 구하면 각 하위 캡슐의 출력이 상위 캡슐들에 분배됩니다. 각 상위 캡슐은 해당 계수로 가중한 예측을 합산한 뒤 스쿼시 함수를 적용합니다. 이 함수는 벡터의 방향을 유지하면서 길이를 0에서 1 사이로 매핑합니다. 각 예측과 상위 캡슐의 새로운 출력 사이의 내적을 로짓에 더하므로, 하위 캡슐은 자신이 예측한 상태를 가진 상위 캡슐에 출력을 더 많이 보냅니다.

import torch

torch.manual_seed(0)

# 하위 캡슐 6개가 상위 캡슐 2개에 대해 8차원 자세를 각각 예측합니다.
predictions = 0.2 * torch.randn(6, 2, 8)
pose = 0.5 * torch.randn(8)
predictions[:4, 0] = pose + 0.05 * torch.randn(4, 8)  # 하위 캡슐 0~3은 상위 캡슐 0과 일치합니다.

logits = torch.zeros(6, 2)
for _ in range(3):  # 논문에 제시된 합의 기반 라우팅의 반복 횟수는 3회입니다.
    coupling = logits.softmax(dim=1)  # 각 하위 캡슐은 표를 상위 캡슐들에 나눠 보냅니다.
    s = (coupling.unsqueeze(-1) * predictions).sum(dim=0)
    squared = s.square().sum(dim=-1, keepdim=True)
    parents = squared / (1 + squared) * s / squared.sqrt()  # 스쿼시: 길이는 [0, 1) 범위입니다.
    logits = logits + (predictions * parents).sum(dim=-1)  # 합의가 로짓을 높입니다.

print(coupling[:, 0].round(decimals=2))
print(parents.norm(dim=-1).round(decimals=2))

tensor([0.9700, 0.9700, 0.9700, 0.9700, 0.5700, 0.4700])과 tensor([0.9800, 0.1500])을 출력합니다. 세 번 반복한 뒤 합의한 하위 캡슐 4개는 출력의 97%를 상위 캡슐 0에 보냅니다. 나머지 2개는 거의 균등하게 분배하며, 산발적인 표만 받은 상위 캡슐 1의 벡터 길이는 0.15인 반면 상위 캡슐 0의 벡터 길이는 0.98에 도달합니다. 라우팅은 Transformer의 어텐션 메커니즘을 반대 방향으로 실행하는 방식과 비슷합니다. 소프트맥스는 하위 캡슐이 출력을 보낼 수 있는 상위 캡슐을 대상으로 실행되므로, 상위 캡슐들이 각 하위 캡슐의 출력을 두고 경쟁합니다.

2017년 CapsNet과 그 성능#

원본 네트워크는 얕은 구조입니다. 크기 9 x 9 커널 256개를 사용하는 합성곱이 PrimaryCaps에 입력되고, 합성곱 캡슐 레이어인 PrimaryCaps는 28 x 28 숫자 이미지를 32 x 6 x 6 = 1,152개의 8차원 캡슐로 변환합니다. 라우팅은 각 클래스에 16차원 캡슐 하나씩을 두는 DigitCaps에 캡슐을 연결합니다. 학습은 캡슐 길이에 마진 손실 함수를 적용해 존재하는 숫자의 길이는 0.9보다 높이고 나머지 숫자의 길이는 0.1보다 낮추며, 올바른 숫자 캡슐로부터 이미지를 복원하는 디코더를 함께 사용합니다. 디코더는 자동 인코더 방식의 정규화 항을 적용하며 가중치는 0.0005로 낮춥니다. 16개 차원 중 하나를 교란한 뒤 결과를 디코딩하면 크기, 획 두께, 기울기, 너비와 같은 속성 하나가 달라집니다. 따라서 캡슐 출력은 CNN 활성화보다 해석하기 쉬우며, 이는 설명 가능한 AI가 추구하는 목표와 같습니다.

기준 CNN의 파라미터가 3,540만 개인 데 비해 파라미터 820만 개를 사용한 CapsNet은 MNIST 테스트 오류율 0.25%를 달성해 기준 모델의 0.39%보다 낮았습니다. 평균적으로 경계 상자가 80% 겹치는 숫자 쌍인 MultiMNIST에서는 8.1% 대비 5.2%를 기록했습니다. 최대 2픽셀 이동한 MNIST 숫자만으로 학습했을 때 무작위 아핀 변환을 적용한 같은 숫자로 구성된 affNIST에서 79%의 점수를 얻었으며, 파라미터 수가 비슷한 CNN은 66%를 기록했습니다. CIFAR-10에서는 CapsNet 7개로 구성된 앙상블이 10.6%의 오류율을 기록했습니다. 이는 CNN을 해당 데이터셋에 처음 적용했을 때의 성능과 비슷한 수준이며, 저자들은 캡슐이 복잡한 배경을 처리하려는 점을 성능 차이의 원인으로 설명했습니다.

한계와 후속 캡슐 모델#

라우팅은 모든 하위 캡슐과 상위 캡슐의 쌍에 대해 여러 차례 반복되므로 이미지 크기와 클래스 수가 커지면 연산량과 메모리 사용량이 빠르게 증가합니다. 캡슐 네트워크는 ImageNet과 같은 대규모 이미지 분류 벤치마크에서 CNN이나 비전 Transformer의 성능을 따라잡지 못했습니다. 연구자들은 예를 들어 소규모 학습 세트를 사용하는 MRI 뇌 스캔의 종양 유형 분류와 같은 의료 영상 분석에 캡슐 네트워크를 적용했습니다.

Hinton, Sabour, Frosst가 ICLR 2018에서 발표한 EM 라우팅을 사용하는 행렬 캡슐은 자세 벡터를 4 x 4 자세 행렬로 대체하고 기대값-최대화 알고리즘으로 라우팅했습니다. 여러 시점에서 촬영한 장난감 물체로 구성된 smallNORB 벤치마크에서 이전 최첨단 모델보다 테스트 오류를 45% 줄였으며, 기준 CNN보다 화이트박스 적대적 공격에 더 강했습니다. 2019년에는 레이블 없이 캡슐을 학습하는 적층형 캡슐 자동 인코더가 등장했습니다. Hinton이 2021년에 제안한 GLOM은 동적 라우팅을 전혀 사용하지 않고, 동일한 부분-전체 계층 구조를 서로 일치하는 임베딩 벡터의 섬으로 표현합니다.