LAYA: Layer-wise Attention Aggregation for Interpretable Depth-Aware Neural Networks

기술 요약: LAYA (Layer-wise Attention Aggregation)

문제 정의
심층 신경망(DNN)은 일반적으로 최종 은닉층($h_L$)이 생성하는 표현에만 전적으로 의존하여 예측을 수행한다. 이러한 설계는 가장 깊은 임베딩이 이전의 변환 과정을 통해 추출된 모든 관련 의미 정보를 완전히 포괄하고 있다는 것을 암묵적으로 가정한다. 그러나 경험적 증거에 따르면 중간층에는 저수준 패턴부터 고수준 추상화에 이르는 풍부하고 상호 보완적인 정보가 포함되어 있다. 표준 아키텍처는 출력 단계에서 이러한 중간 표현들을 버림으로써, 차별적인 정보를 손실할 가능성이 있을 뿐만 아니라, 서로 다른 추상화 수준이 최종 결정에 기여하는 방식을 모호하게 만든다. 더욱이, 이러한 기여도를 해석하기 위한 기존 방법들은 추론 후 외부 분석을 필요로 하는 사후적(post-hoc) 방식에 의존하며, 모델의 예측 메커즘 자체에 내재된 방식이 아니다.

방법론
본 논문은 표준 분류 헤드를 대체하기 위해 설계된 가볍고 아키텍처에 구애받지 않는 출력 모듈인 **LAYA (Layer-wise Attention Aggregator)**를 소개한다. LAYA는 마지막 임베딩만을 투영하는 대신, 모든 은닉층 $\{h_i\}_{i=1}^L$로부터의 내부 표현을 동적으로 집계한다.

메커니즘은 다음과 같이 작동한다:

1. 투영(Projection): 각 은닉 표현 $h_i$는 학습 가능한 어댑터 $g_i(\cdot)$를 통해 공통 잠재 공간(dimension $d^*$)으로 매핑되어 $z_i = g_i(h_i)$를 생성한다.
2. 스코어링(Scoring): 투영된 표현들은 선택적으로 공유 매핑 $\psi(\cdot)$ (항등 함수 또는 경량 MLP)에 의해 변환되고 결합된다. 스코어링 MLP는 이 결합된 값을 처리하여 비정규화된 관련성 점수 $s(x) \in \mathbb{R}^L$를 생성한다.
3. 어텐션(Attention): 점수들은 온도 조절 소프트맥스(temperature-scaled softmax)를 사용하여 입력 조건부 어텐션 계수 $\alpha_i(x)$로 변환된다:
   $$\alpha_i(x) = \frac{\exp(s_i(x)/\tau)}{\sum_{j=1}^L \exp(s_j(x)/\tau)}$$
   여기서 $\tau$는 분포의 선명도(sharpness)를 제어한다.
4. 집계 및 예측(Aggregation & Prediction): 최종 집계된 표현은 가중 합 $h_{agg} = \sum_{i=1}^L \alpha_i(x) z_i$로 계산되며, 이후 태스크 특화 출력 헤드로 전달되어 분류를 수행한다.

결정적으로, 어텐션 가중치 $\alpha_i(x)$는 순전파 과정 중에 예측 프로세스의 일부로서 생성되므로, 추가적인 역전파 계산이나 외부 설명 모델 없이도 내재적 층 기여도 점수(intrinsic layer-attribution scores) 역할을 수행한다.

주요 기여

1. 출력층의 재고: 본 연구는 중간 표현을 무시하는 기존 출력 헤드의 한계를 지적하고, 깊이 인지적 집계(depth-aware aggregation)로의 전환을 제안한다.
2. 내재적 해석 가능성: LAYA는 모델의 결정 과정이 본질적으로 해석 가능하도록 하는 메커니즘을 제공한다. 학습된 어텐션 계수는 각 층이 최종 예측에 기여하는 정도를 명시적으로 정량화하며, 층 단위 기여도 분석을 위해 Integrated Gradients나 Grad-CAM과 같은 사후적 설명 도구를 사용할 필요를 없앤다.
3. 아키텍처 불가지론(Architecture Agnosticism): 이 모듈은 백본 네트워크(예: CNN, Transformer, MLP)를 전혀 변경하지 않고 표준 출력 헤드를 대체할 수 있는 드롭인(drop-in) 교체 방식으로 기능한다.
4. 포괄적 평가: 본 논문은 세 가지 뚜렷한 표현 계층 구조—Fashion-MNIST에 대한 다층 퍼셉트론(MLP), CIFAR-10에 대한 합성곱 신경망(CNN), 예술 양식 분류 작업에 대한 동결된 비전 트랜스포머(ViT)—에 대해 엄격한 평가를 제시한다.

실험 결과

• 예측 성능: LAYA는 표준 베이스라인(LastLayer), 정적 결합(Concat), 글로벌 믹싱(ScalarMix)과 비교하여 경쟁력 있고, 경우에 따라 더 우수한 예측 성능을 달와한다.
  • Fashion-MNIST 및 CIFAR-10에서 LAYA는 테스트된 모든 헤드 중 가장 높은 평균 정확도를 기록했다.
  • Best Artworks 데이터셋에서는 정적 결합(static concatenation)이 가장 좋은 성능을 보였으나, LAYA는 표준 LastLayer 베이스라인을 유의미하게 능가하며 동결된 사전 학습 백본에서도 중간 표현이 가치 있는 정보를 포함하고 있음을 입증했다.
• 설명의 충실도(Faithfulness of Explanations): 기여도 점수와 레이어 제외 섭동 효과(leave-one-layer-out perturbation effects) 사이의 Spearman 상관관계인 Faithfulness와, 가장 높은 순위의 층을 제거했을 때의 신뢰도 하락인 Deletion@1을 사용한 정량적 분석 결과, LAYA의 내재적 점수가 매우 충실함을 보여준다.
  • Fashion-MNIST에서 LAYA의 faithfulness 점수(0.8575)는 경사 기반 사후적 방법(약 0.89)에 근접하면서도 정적 글로벌 가중치보다 훨씬 뛰어난 성능을 보였다.
  • Best Artworks 데이터셋에서 LAYA는 가장 높은 Deletion@1 점수(0.5787)를 달성하였으며, 이는 LAYA가 가장 관련 있다고 식별한 층이 실제로 예측에 가장 큰 영향을 미쳤음을 나타낸다.
• 해석 패턴: 정성적 분석을 통해 태스크 의존적인 깊이 활용 패턴이 밝혀졌다. 예를 들어, CIFAR-10에서는 어텐션이 거의 전적으로 마지막 층에 집중된 반면, Best Artworks에서는 어텐션이 특정 중간 및 후기 ViT 블록(예: 블록 7, 9, 11, 12)에 희소하게 분포되었으며, 이는 예술 양식에 따라 달라졌다.

의의 및 주장
본 논문은 LAYA가 계층적 표현이 출력 단계에서 발생하는 표준적인 "붕괴(collapse)"에 대한 원칙적인 대안을 제공한다고 주장한다. LAYA의 주요 의의는 단순히 분류 정확도를 보편적으로 압도하는 데 있는 것이 아니라, 경쟁력 있는 예측 성능과 직접 검사 가능한 입력 의존적 층 기여도를 결합하는 능력에 있다.

저자는 LAYA가 다음과 같은 기능을 제공한다고 주장한다:

• 내재적 해석 가능성을 제공하여, 외부 도구 없이도 특정 추상화 수준에 대한 모델의 의존성을 투명하게 만든다.
• 구조화되고 태스크 의존적인 깊이 활용 패턴을 드러내며, 서로 다른 태스크와 태스크 내의 서로 다른 클래스들이 서로 다른 조합의 층 깊이에 의존함을 보여준다.
• 어텐션 가이드 모델 압축(어텐션이 미미한 층을 제거하는 프루닝), 적응형 조기 종료(early-exit) 메커니즘, 그리고 어텐션 변화에 기반한 진단적 분석과 같은 미래 응용 분야를 위한 토대를 제공한다.

본 연구는 컴퓨터 비전에 초점을 맞추고 있으나, 이 프레임워크는 다른 시퀀스 아키텍처에도 적용 가능하며, Transformer와 같은 모델에서 깊이에 따라 정보가 어떻게 정제되는지에 대한 새로운 통찰을 제공할 수 있다고 결론짓는다.