[논문 리뷰] LieTransformer: Equivariant self-attention for Lie Groups

Michael Hutchinson, Charline Le Lan외 4명|arXiv (Cornell University)|2020. 12. 20.

Machine Learning in Materials Science참고 문헌 58인용 수 35

한 줄 요약

LieTransformer는 입력을 G 위의 함수로 올려 self-attention을 확장하고, Lie 그룹 및 그 이산 부분군에 대해 등가인 LieSelfAttention 레이어를 구축하여 모양 계수(shape counting), QM9 분자 속성, 해밀토니안 동역학 작업에서 경쟁력 있는 결과를 달성합니다.

ABSTRACT

Group equivariant neural networks are used as building blocks of group invariant neural networks, which have been shown to improve generalisation performance and data efficiency through principled parameter sharing. Such works have mostly focused on group equivariant convolutions, building on the result that group equivariant linear maps are necessarily convolutions. In this work, we extend the scope of the literature to self-attention, that is emerging as a prominent building block of deep learning models. We propose the LieTransformer, an architecture composed of LieSelfAttention layers that are equivariant to arbitrary Lie groups and their discrete subgroups. We demonstrate the generality of our approach by showing experimental results that are competitive to baseline methods on a wide range of tasks: shape counting on point clouds, molecular property regression and modelling particle trajectories under Hamiltonian dynamics.

연구 동기 및 목표

• Lie groups로 설명되는 그룹 대칭을 동기 부여하고 활용하여 학습 효율성과 일반화 성능을 향상시킨다.
• lifting 기반 프레임워크를 통해 self-attention을 Lie groups에서 등가적으로 확장한다.
• 회전/이동 대칭이 있는 작업들에서 접근 방법을 시연한다: 모양 계수(shape counting), 분자 속성 예측, 그리고 해밀토니안 동역학.
• LieTransformer가 Lie 그룹과 그 이산 부분군을 다룰 수 있으며 경쟁력 있는 성능을 보인다.

제안 방법

• 동일 공간 X에서 입력 데이터를 lifting 연산자 L을 통해 Lie 그룹 G 위의 함수로 올려 G-등가 처리를 가능하게 한다.
• LieSelfAttention을 사용하여 리프팅된 도메인 G에서 Content 및 위치 기반 신호를 포함한 주의(attention)를 수행하고, G의 정규 표현 하에서 등가성을 보장한다.
• G 또는 리프팅된 도메인이 연속적일 때 적분을 근사하기 위해 몬테카를로 샘플링을 사용한다 (무한 G_f를 가진 Lie 그룹의 경우).
• LieSelfAttention, LayerNorm, 그리고 MLPs의 잔차 블록을 적용하고, 마지막으로 불변 G-풀링 계층을 통해 작업 출력을 생성한다.
• lifting이 등가이고 LieSelfAttention이 정규 표현 아래에서 등가임을 보임으로써 등가성 증명을 구체화한다.

실험 결과

연구 질문

• RQ1 임의의 Lie 그룹과 그 이산 부분군에 대해 self-attention을 등가적으로 만들 수 있는가?
• RQ2 lifting 기반의 LieSelfAttention 아키텍처가 SE(2)/SE(3) 또는 다른 Lie 그룹 대칭이 필요한 작업에서 경쟁력 있는 성능을 보이는가?
• RQ3 shape counting, 분자 속성 예측(QM9), 및 해밀토니안 다이나믹스 작업에서 LieConv와 비교하여 LieTransformer의 성능은 어떠한가?

주요 결과

• LieTransformer는 shape counting, QM9, 및 해밀토니안 다이나믹스 작업에서 강력한 기준선에 비해 경쟁력 있는 성능을 달성한다.
• LieSelfAttention은 G의 정규 표현 아래에서 증명 가능한 등가성을 가지며(무한 G의 경우 몬테카를로 근사 사용).
• SE(3) 불변 변형은 종종 평행이동만 변형보다 더 나은 성능을 보이고 회전 불변성을 통한 일반화가 향상되며, 리프팅 샘플링으로 인한 일부 분산이 있다.
• 해밀토니안 다이나믹스에서 LieTransformer는 비등가 기준선에 비해 데이터 효율성과 일반화가 현저히 더 우수하다(일부 구간에서 1–3 orders of magnitude 차이가 있다).
• LieConv과 비교할 때 유사한 모델 크기와 그룹에서 LieTransformer가 종종 더 나은 성능을 보이며, 특히 T(2)와 SE(2) 설정에서 두드러진다.