[논문 리뷰] Object Representations as Fixed Points: Training Iterative Refinement Algorithms with Implicit Differentiation

Michael Chang, Thomas L. Griffiths외 1명|arXiv (Cornell University)|2022. 07. 02.

Domain Adaptation and Few-Shot Learning인용 수 5

한 줄 요약

이 논문은 슬롯 어텐션과 같은 반복 보정 모델을 학습하기 위해 백프로파게이션을 전개된 반복에 대해 수행하는 대신 안정적이고 미분 가능한 고정점 공식화로 대체하는 암시적 미분을 제안한다. 이 방법은 일정한 메모리 및 시간 복잡도를 유지하면서도 최적화 안정성 향상, 검증 손실 감소, 생성 품질 향상을 달성하며, 코드 수정은 단 한 줄 뿐이다.

ABSTRACT

Iterative refinement -- start with a random guess, then iteratively improve the guess -- is a useful paradigm for representation learning because it offers a way to break symmetries among equally plausible explanations for the data. This property enables the application of such methods to infer representations of sets of entities, such as objects in physical scenes, structurally resembling clustering algorithms in latent space. However, most prior works differentiate through the unrolled refinement process, which can make optimization challenging. We observe that such methods can be made differentiable by means of the implicit function theorem, and develop an implicit differentiation approach that improves the stability and tractability of training by decoupling the forward and backward passes. This connection enables us to apply advances in optimizing implicit layers to not only improve the optimization of the slot attention module in SLATE, a state-of-the-art method for learning entity representations, but do so with constant space and time complexity in backpropagation and only one additional line of code.

연구 동기 및 목표

• 잠재 집합 표현 학습을 위한 반복 보정 모델에서 전개된 반복에 대해 백프로파게이션을 수행할 경우 발생하는 불안정성과 높은 계산 비용을 해결한다.
• 슬롯 어텐션과 같은 반복 보정 알고리즘의 고정점 성질을 활용하여 암시적 미분을 가능하게 하여 전방 및 역방향 전파를 분리한다.
• 아키텍처 변경 없이 최신 기술인 SLATE 및 일반 슬롯 어텐션 모델의 학습 효율성과 안정성을 향상시킨다.
• 암시적 미분을 통해 기울기 클리핑, 학습률 웜업, 반복 횟수의 초모수 조정 없이도 견고한 학습이 가능함을 입증한다.
• 이 방법이 다양한 아키텍처로 일반화되어 분할 품질을 유지하면서도 후행 작업인 객체 속성 예측 성능을 향상시킴을 보여준다.

제안 방법

• 슬롯 어텐션 모듈을 고정점 반복으로 공식화하여 최종 표현이 z = f_x(z)의 해가 되도록 하여 고정점에서 암시적 미분을 가능하게 한다.
• 암시적 함수 정리를 적용하여 반복 과정을 전개하지 않고 기울기를 계산하며, 효율적인 역방향 전파를 위해 일阶 Neumann 근사법을 사용한다.
• 암시적 기울기를 손실 계산에 통합하여 표준 백프로파게이션과 비교해 단 한 줄의 코드 추가만으로 종단 간 학습이 가능하도록 한다.
• 전방 전파(전개된 반복)와 역방향 전파(암시적 기울기 계산)를 분리하여 각 역방향 전파에서 메모리 및 시간 복잡도를 일정하게 유지한다.
• 보정 함수 f의 야코비안을 사용하여 Neumann 급수 근사법을 통해 암시적 기울기를 계산함으로써 전개된 백프로파게이션으로 인한 기울기 폭발을 방지한다.
• 이 방법을 SLATE 및 원래 슬롯 어텐션 아키텍처에 적용하여 다양한 디코더 및 학습 설정에서의 일반화 성능을 검증한다.

실험 결과

연구 질문

• RQ1 전개된 백프로파게이션으로 인해 기울기 폭발이 발생하는 슬롯 어텐션과 같은 반복 보정 모델의 학습을 암시적 미분이 안정화시킬 수 있는가?
• RQ2 기울기 클리핑 또는 학습률 웜업 없이도 암시적 미분이 객체 중심 표현 학습의 최적화를 향상시킬 수 있는가?
• RQ3 표준 백프로파게이션과 비교했을 때 암시적 미분을 사용할 경우 반복 횟수의 변화가 학습 안정성과 성능에 어떤 영향을 미치는가?
• RQ4 암시적 미분이 다양한 네트워크 아키텍처에서 생성된 마스크 및 분할 출력의 품질을 유지할 수 있는가?
• RQ5 표준 슬롯 어텐션과 비교했을 때 암시적 미분이 객체 속성 예측과 같은 후행 작업에서 얼마나 향상되는가?

주요 결과

• 암시적 슬롯 어텐션은 세 가지 벤치마크 데이터셋에서 표준 슬롯 어텐션 및 SLATE보다 훨씬 낮은 검증 손실을 기록한다.
• SLATE와 비교해 Fréchet Inception Distance(FID)를 최대 30% 감소시키고, 이미지 재구성에서 평균 제곱 오차(MSE)를 최대 50% 감소시킨다.
• 암시적 슬롯 어텐션은 기울기 클리핑, 학습률 웜업, 반복 횟수 조정이 필요 없는 것을 확인했다.
• 역방향 전파에서 메모리 및 시간 복잡도가 반복 횟수에 관계없이 일정하게 유지된다.
• CLEVR 데이터셋에서 일곱 번의 반복을 수행한 암시적 슬롯 어텐션은 단 한 번의 반복만 수행한 표준 슬롯 어텐션보다 재구성 MSE가 2배 낮다.
• 암시적 슬롯 어텐션은 원래 Locatello 등 [39]의 공간 브로드캐스트 디코더를 포함한 다양한 아키텍처에서도 높은 품질의 분할 마스크를 유지한다. 그림 8에서 이를 확인할 수 있다.