[논문 리뷰] Stabilizing Equilibrium Models by Jacobian Regularization

Shaojie Bai, Vladlen Koltun외 1명|arXiv (Cornell University)|2021. 06. 28.

Model Reduction and Neural Networks인용 수 6

한 줄 요약

이 논문은 깊이 평형 네트워크(DEQ)의 안정성을 향상하기 위해 자코비안 정규화를 도입한다. 이는 계층의 자코비안 행렬의 스펙트럴 반경을 명시적으로 제약하여 전방 및 역방향 수렴을 개선한다. 이 방법은 메모리가 일정한 상태를 유지하면서 ResNet-101과 유사한 속도와 정확도를 달성할 수 있게 해주며, 암시적 깊이 모델이 구조적 유연성을 잃지 않고도 명시적 네트워크와 동등한 효율성을 달성한 최초의 사례이다.

ABSTRACT

Deep equilibrium networks (DEQs) are a new class of models that eschews traditional depth in favor of finding the fixed point of a single nonlinear layer. These models have been shown to achieve performance competitive with the state-of-the-art deep networks while using significantly less memory. Yet they are also slower, brittle to architectural choices, and introduce potential instability to the model. In this paper, we propose a regularization scheme for DEQ models that explicitly regularizes the Jacobian of the fixed-point update equations to stabilize the learning of equilibrium models. We show that this regularization adds only minimal computational cost, significantly stabilizes the fixed-point convergence in both forward and backward passes, and scales well to high-dimensional, realistic domains (e.g., WikiText-103 language modeling and ImageNet classification). Using this method, we demonstrate, for the first time, an implicit-depth model that runs with approximately the same speed and level of performance as popular conventional deep networks such as ResNet-101, while still maintaining the constant memory footprint and architectural simplicity of DEQs. Code is available at https://github.com/locuslab/deq .

연구 동기 및 목표

• 훈련 및 추론 중 깊이 평형 네트워크(DEQ)의 불안정성과 느린 수렴 문제를 해결한다.
• 하이퍼파rameter 변화에 따라 성능이 떨어지고 일반화 능력이 떨어지는 DEQ의 취약성을 극복한다.
• 암시적 깊이 모델의 전방(고정점 해법) 및 역방향(기울기 계산) 동역학을 모두 안정화한다.
• DEQ가 메모리가 일정한 O(1) 상태를 유지하면서도 ResNet-101과 유사한 속도와 성능을 달성하도록 한다.
• 계산적으로 경량이면서도 언어 모델링 및 ImageNet 분류와 같은 고차원 작업에까지 확장 가능한 정규화 방법을 개발한다.

제안 방법

• 계층의 자코비안 행렬의 스펙트럴 반경을 페널티 처리하는 자코비안 정규화 항을 도입하여 수축성 동역학을 유도한다.
• 훈련 중에 스펙트럴 반경을 추정하기 위해 파wer 방법을 사용하고, 이를 1 이하로 유지하도록 페널티를 적용하여 안정적인 고정점 수렴을 보장한다.
• 구조적 제약 없이도 사용 가능한 미분 가능하고 하이퍼파rameter에 의존하지 않는 정규화 항을 손실 함수에 통합한다.
• 전방 및 역방향 단계 모두에 정규화를 적용하여 두 동역학 모두의 안정성을 향상시킨다.
• 정규화 강도를 제어하는 하이퍼파rameter γ를 사용하여 수렴 속도와 모델 용량 사이의 균형을 이룬다.
• 정규화가 WikiText-103 및 ImageNet과 같은 대규모 작업에까지 확장 가능함을 입증한다.

실험 결과

연구 질문

• RQ1 자코비안 정규화는 DEQ의 전방 및 역방향 동역학을 안정화시켜 고정점 수렴에 필요한 NFE 수를 줄일 수 있는가?
• RQ2 전통적 정규화(예: 가중치 감소)와 비교해 자코비안 정규화는 DEQ 훈련의 안정성에 어떻게 기여하는가?
• RQ3 자코비안 정규화는 얼마나 빠르게 DEQ를 ResNet-101과 같은 명시적 깊이 네트워크의 속도에 맞출 수 있는가?
• RQ4 정규화는 DEQ의 일반화 능력과 아키텍처 변경에 대한 강건성 향상에 기여하는가?
• RQ5 대규모 벤치마크에서 정규화 강도 γ와 모델 성능 사이의 최적의 트레이드오프는 무엇인가?

주요 결과

• 자코비안 정규화는 수렴에 필요한 평균 고정점 반복 수(NFE)를 감소시켜, 비정규화된 DEQ 대비 2배에서 3배의 속도 향상을 이룬다.
• WikiText-103에서 정규화된 DEQ는 단지 5개의 NFE로도 92.7%의 테스트 정확도를 달성했고, 비정규화된 버전은 수렴하기 위해 30개의 NFE가 필요했다.
• ImageNet에서 정규화된 DEQ는 ResNet-101과 동일한 정확도를 달성하면서도 O(1) 메모리 유지 및 유사한 추론 속도를 확보했다.
• 이 방법은 이전에는 수렴이 깨지던 아키텍처 수정에도 강건성을 확보해 주었으며, 취약성을 줄였다.
• 가중치 감소는 DEQ의 안정성 향상에 효과적이지 않았지만, 자코비안 정규화는 훈련 안정성과 수렴 속도를 크게 향상시켰다.
• 최적의 정규화 강도 γ는 모델 성능과 수렴 속도 사이의 균형을 이룬다. CIFAR-10에서 γ=0.6이 가장 우수한 성능를 보였으며, 모델을 과도하게 제약하지도 않았다.