등변 신경망 3부: 텐서 필드 네트워크와 SE(3) 트랜스포머로 보는 3D 회전 대칭

핵심 메시지

• 트랜스포머와 GNN은 입력 순서가 바뀌어도 결과가 유지되는 순열 등변성은 갖지만, 3차원 회전에 대한 등변성은 갖고 있지 않다.
• 3차원 회전은 2차원과 달리 순서를 바꾸면 결과가 달라지기 때문에, 회전 등변 신경망을 설계하는 일이 훨씬 까다로워진다.
• SchNet은 점 사이의 거리만 사용해 회전에 영향받지 않지만, 방향 정보를 잃어 거울상 구분이나 방향성 힘 예측에는 쓸 수 없다.
• 텐서 필드 네트워크는 구면 조화함수·위그너 D 행렬·클렙슈-고르단 계수를 조합해 회전에 예측 가능하게 반응하는 필터를 만든다.
• SE(3) 트랜스포머는 어텐션의 가중치 행렬을 이 등변 커널로 갈아끼워, 회전과 평행이동에 등변인 어텐션을 구현한다.

쉽게 이해하기

이 영상은 등변 딥러닝 시리즈의 마지막 편으로, 집합·그래프·점군 데이터를 다루는 트랜스포머와 그래프 신경망에 대칭성을 어떻게 심을 수 있는지 다룬다. 점군과 그래프는 모두 '데이터 점들의 모음'인 집합의 일종이고, 여기에 구조를 주면 그래프가, 좌표를 더하면 점군이 된다. 트랜스포머는 CNN 같은 강한 귀납 편향이 없어 다양한 데이터에 유연하게 적용되고, GNN은 메시지 전달을 위해 그래프 구조를 요구한다. 두 모델 모두 입력 순서에 둔감한 순열 등변성은 이미 갖고 있지만, 회전에 대해서는 그렇지 않다는 것이 출발점이다.

목표는 모델이 절대 좌표를 무시하고 점들 사이의 상대적인 정보를 쓰게 만드는 것이다. 입력을 회전시키면 출력도 그만큼 회전하는 성질을 갖게 하면, 모델이 전역 회전의 효과를 따로 학습할 필요가 없어 학습이 단순해진다. 발표자는 3차원 회전이 교환법칙을 따르지 않는다는 점을 상자와 화살표 그림으로 보여준다. y축으로 45도 돌린 뒤 z축으로 180도 돌린 결과와, 순서를 바꾼 결과는 화살표 방향이 서로 반대가 된다.

초기 접근인 SchNet은 격자 대신 모든 점 위에서 정의되는 연속 필터를 쓰고, 점들 사이의 쌍별 거리를 활용한다. 거리는 방향을 바꿔도 변하지 않으므로 이 모델은 회전 불변이지만 등변은 아니다. 그 대가로 벡터의 방향 정보가 사라져, 거울상으로 뒤집힌 물체를 구분하거나 방향이 있는 힘을 예측하는 일이 불가능해진다.

텐서 필드 네트워크는 SchNet의 연속 필터에, 이전 편에서 다룬 조화 기저함수 아이디어를 결합한다. 2차원의 원형 조화함수에 대응하는 3차원 도구가 구면 조화함수이고, 이들이 정규직교 기저를 이루기 때문에 구 위의 임의의 함수를 이들의 합으로 표현할 수 있다. 특징은 회전 불변 스칼라인 0형과 3차원 좌표인 1형처럼 유형으로 나뉘고, 각 유형이 회전에서 어떻게 변환되는지는 SO(3)의 기약 표현인 위그너 D 행렬이 알려준다. 서로 다른 유형의 텐서를 곱할 때 어떤 성분끼리 결합해야 하는지는 클렙슈-고르단 계수가 정해 준다.

마지막으로 다루는 SE(3) 트랜스포머는 텐서 필드 네트워크 위에 세워진 모델이다. 먼저 각 점을 중심으로 일정 반경 안의 이웃만 모아 국소 그래프를 만들어, 모든 점이 서로를 참조할 때 생기는 제곱 복잡도를 피한다. 그다음 쿼리·키·값을 만드는 일반적인 가중치 행렬 자리에 구면 조화함수와 학습 가능한 반경 함수, 클렙슈-고르단 계수로 구성된 등변 커널을 넣는다. 어텐션 점수 자체는 내적과 소프트맥스로 평소와 똑같이 계산되지만, 이 교체만으로 어텐션 전체가 3차원 회전과 평행이동에 등변이 된다.

주요 인사이트

• 불변과 등변은 다르다. 거리만 쓰면 회전에 흔들리지 않는 대신 방향 정보를 통째로 버리게 되므로, 힘이나 방향을 예측해야 하는 문제에서는 등변 설계가 필요하다.
• 대칭성을 모델에 넣는 것은 표현력을 깎는 제약이 아니라 학습 부담을 덜어주는 장치다. 전역 회전이 특징에 미치는 영향을 데이터로 배우지 않아도 되기 때문이다.
• 3차원 회전이 교환법칙을 따르지 않는다는 성질이 설계 전반을 제약한다. 2차원에서 쓰던 조화함수 기법을 그대로 옮길 수 없고, 텐서 유형별 결합 규칙이 필요한 이유가 여기에 있다.
• SE(3) 트랜스포머의 핵심 아이디어는 새로운 어텐션을 발명한 것이 아니라, 기존 어텐션의 가중치 행렬을 등변 커널로 바꿔 끼운 것이다. 익숙한 구조를 유지한 채 대칭성만 주입한 셈이다.
• 발표자는 이 시리즈에서 다루지 못한 모델이 많지만, 정리한 개념만으로도 대부분의 등변 모델을 읽어낼 수 있다고 말한다. 실습용 노트북과 저자 구현, 계산 속도를 크게 높인 엔비디아 구현이 함께 소개된다.

자주 묻는 질문

관련 AI 소식