Enhancing Multiple Reliability Measures via Nuisance-extended Information Bottleneck
논문의 핵심 요약
이 논문은 딥러닝 모델이 학습 데이터에만 존재하는 shortcut에 과도하게 의존하여, 실제 환경에서 예측이 불안정해지는 문제를 해결하고자 합니다. 예를 들어, 소 사진을 학습할 때 푸른 초원 배경과 함께 학습하면, 모델은 '소' 자체보다 '푸른 초원'을 보고 소라고 판단하는 실수를 저지를 수 있습니다.
이 문제를 해결하기 위해 저자들은 '정보 병목(Information Bottleneck, IB)' 원칙을 확장하여 '불필요 정보(Nuisance Information)' 라는 개념을 도입합니다. 모델이 이미지를 분류하는 데 꼭 필요한 핵심 정보(z)와, 분류에는 불필요하지만 이미지를 구성하는 나머지 정보(z_n, 예: 배경, 조명, 질감 등)를 명확히 분리하도록 학습시키는 새로운 프레임워크 NIB(Nuisance-extended Information Bottleneck)를 제안합니다.
0. Abstract
실제 시나리오에서 training data가 제한적인 경우, 데이터 내 많은 predictive signal은 데이터 수집 과정의 bias에서 기인하는 경우가 많음(즉, generalizable하지 않음).
--> training dist밖에선 안통함 bias에서 온 이상한 학습된거란말
이로 인해 모델이 이러한 (소위) “shortcut” signal에 co-adaptation 되는 것을 막기 어려움. 이는 다양한 distribution shift 상황에서 모델을 fragile하게 만듦.
co-adaptation = 여러 feature(또는 neuron)들이 서로에게 의존하면서 특정 패턴에 과하게 맞춰 학습되는 현상.
이러한 failure mode를 우회하기 위해, training 과정에서 더 넓은 perturbation class를 포괄할 수 있도록 mutual information constraint 하의 adversarial threat model을 고려함.
perturbation class = 이런 변형들의 집합 (adversarial noise (ℓp-ball 안 노이즈) fog, brightness (corruption) background 변경 style 변화 카메라 노이즈... 등 특정 perturbation들만 학습)
-> 어떤 형태든 y의미가 유지되고 이외 다른 관련없는 변화들은 다 포함하고싶음
이에 따라 standard Information Bottleneck(IB)를 확장하여 nuisance information을 추가적으로 모델링하도록 설계함.
해당 objective를 구현하기 위해 autoencoder 기반 training을 제안함. 또한 convolutional 기반 및 Transformer 기반 아키텍처 모두에서 제안하는 hybrid discriminative-generative training을 원활히 수행할 수 있도록 practical encoder design을 함께 제안함.
실험 결과, 제안한 방법은 learned representation의 robustness를 향상시킴이 확인됨(특히 어떠한 domain-specific knowledge도 사용하지 않았음에도 불구함). 이는 여러 challenging reliability measure에 대해 일관되게 나타남.
예를 들어, 본 모델은 최근 challenging한 OBJECTS benchmark의 novelty detection task에서 AUROC 기준 78.4% → 87.2%로 state-of-the-art를 달성함.
동시에 corruption robustness, background robustness, 그리고 (certified) adversarial robustness 또한 향상됨.
코드는 https://github.com/jh-jeong/nuisance_ib 에 공개되어 있음.
OBJECTS benchmark :
기존 OOD benchmark는 너무 다른 데이터로 테스트해서 쉬운 편임.
OBJECTS는 CIFAR-10 + CIFAR-10-C + ImageNet-10을 모두 in-distribution으로 두고,
그와 의미적으로 가까운 데이터까지 구분해야 하도록 만든 더 어려운 OOD 평가 세팅임.
즉, semantic하게 애매한 경우까지 구분해야 하는 강화된 OOD benchmark
1. Introduction
deep learning의 도움으로 image/video recognition [9, 20, 109], synthesis [42, 55, 96, 125], 3D scene rendering [85, 86, 104] 등 computer vision 분야에서 최근 큰 breakthrough가 있었음에도 불구하고, deep learning 모델을 real-world에 배포하는 일은 여전히 contents provider에게 부담이 됨.
이는 서비스의 reliability에 직접적인 영향을 미치기 때문임.
많은 경우, deep neural network는 out-of-distribution(OOD) 입력, 즉 training distribution에서 나왔을 가능성이 낮은 sample에 대해 상당히 fragile한 예측을 수행함.
심지어 해당 입력이 인간의 관점에서는 in-distribution sample과 semantic하게 충분히 유사한 경우에도 그러함 [35, 102].
이러한 취약성은 autonomous driving, medical imaging, health-care application 등과 같은 risk-sensitive system에서 심각한 위협이 될 수 있음 [2].
전반적으로 이러한 현상은 deep neural network가 실제 환경에서 제한적이거나(또는 잠재적으로 biased된) training data로부터 “shortcut” signal [26]을 추출하는 경향이 있음을 보여줌.
이러한 문제를 해결하기 위해, model reliability의 서로 다른 측면에 기반한 다양한 연구들이 독립적으로 발전되어 왔음.
이들 방법은 test-time에서 input distribution shift가 어떻게 발생하는지, 그리고 그 shift에 대해 모델 성능을 어떻게 평가하는지에 따라 서로 다른 threat model과 benchmark를 사용함.
- 예를 들어 adversarial robustness [11, 17, 82, 128] 맥락에서는, 주어진 test sample 주변의 고정된 ℓp-ball 내부에서 worst-case noise를 고려하는 것이 전형적인 threat model임.
- 또 다른 예로 corruption robustness [34, 35, 39, 116]에서는 Gaussian noise, fog 등과 같은 사전에 정의된 common corruption이 test sample에 적용된다고 가정함.
- 마지막으로 novelty detection [36, 72, 74, 76]에서는 특정 benchmark dataset을 (in-distribution) test sample과 구분하여 out-of-distribution으로 탐지할 수 있는지를 평가함.
여기서 reliability =
분포가 바뀌거나 (distribution shift) 입력이 변형되더라도모델이 계속 신뢰할 수 있게 동작하는 성질
- accuracy만 높은 게 아니라
- OOD에서도 무너지지 않고
- corruption에도 견디고
- adversarial에도 덜 깨지고
- uncertainty도 제대로 표현하는 것
robustness ⊂ reliability
Robustness는 특정 공격에 안 깨지는 성질이고,
Reliability는 전반적으로 믿을 수 있는 성질임.
모델이 distribution shift(테스트 데이터가 훈련 데이터와 달라지는 현상)에 약하다는 문제가 있어서,
이를 해결하려는 연구들이 많이 나왔는데,
👉 각 연구마다 “문제가 생기는 방식”을 다르게 정의함
👉 그래서 쓰는 threat model(입력이 어떤 방식으로 망가질 수 있다고 가정할 것인가?를 정의한 가정)
이 다르고, benchmark도 다름
--> 테스트 데이터가 어떻게 달라질 거라고 가정하느냐 에 따라 연구 방향이 달라졌다는 말임.
1️⃣ Adversarial robustness 연구
- 가정: 입력에 아주 작은 ℓp 노이즈가 붙음
- threat model: ℓ∞ ball 안의 worst-case noise
- benchmark: PGD attack 등
- 특정 범위 내 공격에 대해 수학적 안전 보장
2️⃣ Corruption robustness 연구
- 가정: fog, blur, noise 같은 자연스러운 변형 발생
- threat model: 미리 정의된 corruption 종류
- benchmark: CIFAR-10-C, ImageNet-C
정의
자연스러운 이미지 변형이나 노이즈가 발생해도 성능이 유지되는 능력이다.
변형 예시
- blur
- fog
- brightness 변화
- Gaussian noise
- JPEG compression
핵심 질문
“입력에 자연스러운 왜곡이 생겨도 모델이 여전히 잘 맞추는가?”
의미
현실 세계에서는 항상 깨끗한 데이터만 들어오지 않는다.
따라서 실제 환경에서의 일반화 능력을 평가하는 지표이다.
3️⃣ Novelty detection 연구 (OOD Detection)
- 가정: 완전히 다른 데이터셋이 들어옴
- threat model: OOD dataset
- benchmark: SVHN, LSUN, OBJECTS
정의
훈련 데이터 분포에 속하지 않는 새로운 입력(out-of-distribution, OOD)을 탐지하는 능력이다.
핵심 질문
“이 입력이 내가 학습한 데이터와 같은 종류인가, 아니면 완전히 다른 것인가?”
예시
- CIFAR-10으로 학습한 모델에 MNIST가 들어옴
- 정상 의료 영상으로 학습했는데 전혀 다른 장비에서 찍은 영상이 들어옴
중요성
모델이 모르는 데이터를 자신 있게 잘못 예측하는 것은 매우 위험하다.
Novelty detection은 모델이 “모른다”고 말할 수 있게 하는 능력이다.
4. Background Robustness
정의
객체의 본질적인 특징이 아니라 배경 정보에 의존하지 않는 능력이다.
그러나 각 “ideal” objective와 실제 threat model 사이의 불일치로 인해, 특정 threat model 하에서 최적화한 방법은 다른 threat model로 잘 일반화되지 않는다는 점이 공통적으로 관찰됨.
- 예를 들어, (a) 여러 연구 [16, 62, 121]에서는 standard adversarial training [82]이 corruption robustness나 uncertainty estimation과 같은 다른 reliability measure뿐 아니라 classification performance까지도 저해하는 경우가 많음을 보고함.
- (b) Hendrycks et al. [34]는 기존 corruption robustness에 대한 주장들이 보다 종합적인 benchmark에서는 일관되게 일반화되지 못함을 비판함. 이러한 현상은 동일한 objective를 목표로 하는 threat model 간에도 발생함.
- 예를 들어, (c) Yang et al. [122]는 novelty detection 분야의 state-of-the-art 방법들이 지나치게 민감하여 “near-in-distribution” sample까지 out-of-distribution으로 탐지하는 경향이 있으며, 이를 평가하는 benchmark에서 성능이 저하됨을 보임.
- 전반적으로 이러한 관찰은 특정 threat model이나 benchmark를 가정한 채 reliability measure를 최적화하는 접근을 지양해야 함을 시사하며, 다양한 reliability 시나리오에 일반적으로 적용 가능한 새로운 threat model의 필요성을 동기 부여함.
Contribution.
본 논문에서는 domain-specific task에 대한 prior를 가정하지 않고 model reliability를 목표로 하는 새로운 training objective인 nuisance-extended information bottleneck(NIB)을 제안함.
본 방법은 distribution shift가 존재하는 상황에서 information bottleneck(IB) 원리 [107, 108]를 재고하는 데서 출발함.
구체적으로, “robust” representation \(z := f(x)\)는 입력 ‘\(x\)’에서 target ‘\(y\)’와 상관관계가 있는 모든 signal을 인코딩해야 하며, 일부 shortcut만을 추출해서는 안 된다고 주장함(예: Figure 1).
이에 따라 입력 ‘\(x\)’의 distribution shift에 대해, mutual information ‘\(I(x, y)\)’ 제약 하에서 adversarial한 형태의 threat model을 고려함.
이 아이디어를 구현하기 위해, standard IB의 representation ‘\(z\)’와 함께 nuisance representation ‘\(z_n\)’을 도입하여 ‘\((z, z_n)\)’이 ‘\(x\)’를 재구성할 수 있도록 하는 practical design을 제안함.
이는 adversarial autoencoder [83]와 variational IB [1]를 하나의 framework로 통합한 새로운 synthesis로 이어짐.
아키텍처 측면에서는,
(a) convolutional network 기반 encoder에 대해 internal feature statistic을 활용하고,
(b) Transformer 기반 [24] encoder에 대해 vector-quantized patch representation을 도입하여 ‘\(z_n\)’을 모델링함.
이는 nuisance ‘\(z_n\)’(및 ‘\(z\)’ 자체)을 scalable하게 효율적으로 인코딩하기 위함임.
제안한 방법으로 학습된 representation에 대해 extensive evaluation을 수행하였으며, 최신 reliability metric 전반에 걸쳐 포괄적인 개선을 보임.
여기에는 (a) novelty detection, (b) corruption(또는 natural) robustness, (c) background robustness, (d) certified adversarial robustness가 포함됨.
특히 이러한 성능 향상은 각 threat model에 대한 domain-specific prior를 가정한 결과가 아니라는 점에서 의미가 큼.
예를 들어, 최근 방법들 [39, 40]에서 가정한 domain-specific prior 없이도 CIFAR-10-C의 최고 severity error rate를 26.5% → 19.5%로 유의미하게 감소시킴.
또한 본 방법의 효과가 대규모(ImageNet) dataset으로도 확장 가능함을 보임.
novelty detection에서는 최근 OBJECTS [122] benchmark에서 AUROC을 평균적으로 78.4% → 87.2%까지 크게 향상시켜 state-of-the-art를 넘어섰으며, 이는 제안한 representation이 OOD sample을 더 잘 구분할 수 있는 semantic 정보를 제공함을 의미함.
마지막으로, 학습된 representation에 randomized smoothing [17]을 적용함으로써 adversarial example에 대해서도 향상된 robustness를 제공할 수 있음을 보임.
Randomized smoothing은 분류 모델에 대해 수학적으로 보장된(certified) adversarial robustness를 제공하는 방법이다. 기본 아이디어는 하나의 입력에 대해 단일 예측을 사용하는 대신, 해당 입력에 **가우시안 노이즈(Gaussian noise)**를 반복적으로 추가한 뒤 예측을 평균내어 최종 결정을 내리는 것이다.
구체적으로, 원래 분류기 f(x)가 있을 때, 새로운 분류기 g(x)를 다음과 같이 정의한다.
g(x)=argmaxcP(f(x+ϵ)=c)g(x) = \arg\max_c \mathbb{P}(f(x + \epsilon) = c)
여기서
ϵ∼N(0,σ2I)\epsilon \sim \mathcal{N}(0, \sigma^2 I)
즉, 입력 xx에 평균이 0이고 분산이 σ2\sigma^2인 가우시안 노이즈를 반복적으로 더한 뒤, 가장 높은 확률로 예측되는 클래스를 최종 출력으로 선택한다.
이 방식의 핵심은 다음과 같다. 만약 입력 xx 주변에 작은 변형(perturbation)이 가해지더라도 모델의 예측이 대부분 동일하게 유지된다면, 해당 입력 주변은 안정적이라고 볼 수 있다. Randomized smoothing은 이러한 안정성을 이론적으로 분석하여, 특정 반경 rr 이내의 모든 입력 변형에 대해 예측이 변하지 않는다는 것을 수학적으로 보장한다. 이때의 rr을 certified radius라고 하며, 이 값이 클수록 더 강한 adversarial robustness를 의미한다.
기존 adversarial training은 경험적으로 공격에 강한 모델을 만들 수 있지만, “어느 범위까지 안전한지”에 대한 명확한 보장을 제공하지 못한다. 반면 randomized smoothing은 확률적 분석을 통해 명시적인 안전 반경을 계산할 수 있다는 점에서 차별화된다.
해당 논문에서는 nuisance modeling을 통해 학습된 표현이 입력 변형에 대해 보다 안정적으로 작동함을 보이고, 여기에 randomized smoothing을 적용함으로써 별도의 복잡한 adversarial training 없이도 높은 certified adversarial robustness를 달성할 수 있음을 실험적으로 확인하였다.
요약하면, randomized smoothing은 입력에 가우시안 노이즈를 반복적으로 적용하여 평균 예측을 사용하는 기법이며, 이를 통해 일정 반경 내의 adversarial perturbation에 대해 예측이 변하지 않음을 이론적으로 보장하는 방법이다.
mutual information이 뭐냐?
I(x;y)I(x; y)
= x가 y에 대해 얼마나 많은 정보를 가지고 있는지
쉽게 말하면:
- x를 보면 y를 얼마나 알 수 있냐?
- x와 y의 통계적 의존성의 크기
예시
- x = 강아지 사진
- y = “강아지”
→ mutual information 큼
- x = 완전 랜덤 노이즈
- y = “강아지”
→ mutual information 거의 0
논문에서 말하는 가정
논문은 이렇게 가정함:
I(x;y)=I(x^;y)I(x; y) = I(\hat{x}; y)
즉,
입력이 변형되어도 의미(y에 대한 정보)는 그대로라고 가정함
예:
- 색깔 바뀜
- 배경 바뀜
- 약간 blur
하지만, “개”는 여전히 “개”
그럼 adversarial threat model ?
논문은 말함:
“그럼 의미는 그대로인데 모델을 가장 힘들게 만드는 입력 변형을 가정하자”
즉,
minx^I(f(x^);y)\min_{\hat{x}} I(f(\hat{x}); y)
의미는 유지되는데 모델이 y를 못 맞추게 만드는 방향으로 입력이 변한다고 상상함.
이게 adversarial 형태임.
vector-quantized patch representation = 이미지 patch embedding을 코드북 기반으로 이산화한 표현
→ generative하고 안정적인 nuisance modeling을 위해 사용됨.
2. Background
Notation.
두 random variable ‘\(x \in \mathcal{X}\)’(input)와 ‘\(y \in \mathcal{Y}\)’(target)가 주어졌을 때, 우리의 목표는 데이터 ‘\(D = \{(x_i, y_i)\}_{i=1}^n\)’로부터 mapping(또는 encoder) ‘\(f : \mathcal{X} \to \mathcal{Z}\)’를 학습하는 것임. 여기서 representation ‘\(z := f(x)\)’는 더 단순한(e.g., linear) mapping [7, 64]을 통해 ‘\(y\)’를 예측할 수 있어야 함. ‘\(f\)’는 neural network로 parametrization되었다고 가정하며, information-theoretic 관점을 채택하기 위해 stochastic하게 정의함 [108]. 즉, encoder output은 상수가 아니라 ‘\(p_f(z|x)\)’로 정의되는 random variable임. 이러한 모델링은 independent random variable ‘\(\epsilon\)’과 (deterministic) ‘\(f\)’를 이용한 reparametrization trick [61]을 통해 구현 가능하며, ‘\(z := f(x, \epsilon)\)’로 가정함. 예를 들어, 표준 설계 중 하나는 다음과 같이 ‘\(f\)’를 parametrization함:
‘\(f(x, \epsilon) := f_\mu(x) + \epsilon \cdot f_\sigma(x)\)’ (1)
여기서 ‘\(f_\mu \in \mathbb{R}^{|\mathcal{Z}|}\)’ 및 ‘\(f_\sigma \in \mathbb{R}^{|\mathcal{Z}|}_{+}\)’는 각각 ‘\(\mathcal{N}(x; \mu, \sigma^2 I)\)’에서 ‘\(\mu\)’와 ‘\(\sigma\)’를 모델링하는 deterministic mapping이며, gradient-based optimization을 통해 학습 가능함.
데이터 ‘\(D\)’는 일반적으로 어떤 data-generating distribution으로부터의 i.i.d. sample ‘\((x_i, y_i) \sim p_d(x, y)\)’로 구성된다고 가정함. ‘\(D\)’로부터 학습된 ‘\(f\)’는 보지 못한 sample에 대해 ‘\(p_d(y|x)\)’를 잘 generalize하리라 기대함. 그러나 이 formulation은 ‘\(p_d\)’에서 나올 가능성이 낮은 입력, 예컨대 ‘\(\hat{x}\)’에 대해 ‘\(f\)’가 어떻게 동작해야 하는지 명시하지 않음. 이는 특히 인간이 nuisance로 간주하는 변화만 존재하는 경우에도 ‘\(f\)’의 decision making이 인간과 유사해야 한다고 기대하는 상황에서 문제가 됨. 이러한 지점이 standard training practice 하에서 current neural network가 흔히 실패하는 부분임.
Information bottleneck.
직관적으로 “좋은” representation ‘\(z\)’는 ‘\(x\)’ 중 ‘\(y\)’와 상관된 정보를 유지하면서도, 지나치게 복잡하지 않아야 함. Information Bottleneck(IB) [107, 108]은 주어진 downstream task ‘\(x \to y\)’에 대해 간결한 representation ‘\(x \to z\)’를 얻기 위한 principled한 접근임. 즉, (a) task-relevant mutual information ‘\(I(z; y)\)’를 최대화하면서, (b) better generalization을 위해 ‘\(z\)’의 capacity를 제한하기 위해 ‘\(I(x; z)\)’를 최소화함. 다시 말해, ‘\(I(x; z)\)’를 ‘\(z\)’의 complexity measure로 설정함. 구체적으로 다음 objective를 최대화함:
‘\(\max_f R_{IB}(f)\), for \(R_{IB}(f) := I(z; y) - \beta I(x; z)\)’ (2)
여기서 ‘\(\beta \ge 0\)’는 capacity constraint를 제어하는 계수이며, 어떤 ‘\(I_\beta\)’에 대해 ‘\(I(x; z) \le I_\beta\)’가 되도록 보장함.
IB에서는 I(x;z)를 줄이고 I(z;y)를 늘리려함
- 를 키우는 이유 : z가 y에 대해 얼마나 많은 정보를 담고 있는가
- z를 보면 y를 잘 맞출 수 있어야 함 예측에 필요한 정보는 유지해야 함 그래서 최대화함.
- 를 줄이는 이유: z가 입력 x를 얼마나 많이 기억하고 있는가
- 값이 크면: z가 x를 거의 그대로 복사함 배경, 노이즈, shortcut까지 다 들어감
과적합 위험 증가 distribution shift에서 깨짐
그래서: z를 “압축”해서 진짜 중요한 정보만 남기게 해야 함 → 최소화
- 값이 크면: z가 x를 거의 그대로 복사함 배경, 노이즈, shortcut까지 다 들어감
3. Nuisance-extended information bottleneck
standard Information Bottleneck(IB) objective (2)는 향후 입력이 또한 ‘\(p_d(x, y)\)’로부터 생성된다는 전제 하에 representation ‘\(z := f(x)\)’를 학습함.
본 논문에서는 미래에 입력 ‘\(x\)’가 어떤 unknown noisy channel을 통해 corrupt되어 ‘\(x \to \hat{x}\)’가 될 수 있다고 가정함.
단, ‘\(\hat{x}\)’는 ‘\(y\)’에 대한 semantic을 보존한다고 가정함.
즉, ‘\(I(x; y) = I(\hat{x}; y) > 0\)’를 가정함.
직관적으로, ‘\(x\)’가 각각 이미 ‘\(y\)’와 높은 상관관계를 갖는 multiple signal을 포함하고 있는 상황을 상상할 수 있음.
즉, 일부를 제거해도 ‘\(y\)’와의 mutual information에는 영향을 주지 않는 경우임.
이러한 signal은 deep neural network에서 실제로 흔히 존재함. 예를 들어 [44]는 adversarial perturbation [29, 102]만으로도 정확한 classification이 가능함을 경험적으로 관찰함.
IB framework 맥락에서, 간결한 encoder \(f\)를 얻는 것이 목표라면, noisy channel \(\hat{x}\)는 adversary처럼 작동한다고 보는 것이 합리적임. 즉, 다음을 최소화하는 adversarial setting을 고려함:
‘\(\min_{\hat{x}} I(\hat{z} := f(\hat{x}); y)\)’ subject to ‘\(I(x; y) = I(\hat{x}; y)\)’ (3)
이는 channel의 사전적 거동에 대한 정보가 없다는 가정 하에서 설정됨. 이러한 최적화는 ‘\(x\)’에서 ‘\(y\)’와 높은 상관관계를 갖는 모든 signal을 ‘\(f\)’가 반드시 추출하도록 요구함. 그렇지 않으면, ‘\(\hat{x}\)’가 ‘\(f\)’가 놓친 단 하나의 signal만 남기고 나머지를 모두 제거하는 경우가 발생할 수 있음. 그러나 실제로는 (a) 고차원 공간 ‘\(\mathcal{X}\)’에서의 사실상 unconstrained optimization이며, (b) 계산하기 어려운 mutual information 제약을 포함한다는 점에서, (3)을 직접 최적화하는 것은 computationally infeasible함.
본 논문에서는 (3)을 직접 사용하지 않으면서도 \(f\)가 adversarial behavior를 보이도록 하기 위해, \(f\)가 representation \(z\)뿐 아니라 nuisance representation \(z_n\)도 함께 모델링하도록 제안함.
구체적으로, ‘\(z_n\)’은 ‘\(x\)’를 재구성하기 위해 ‘\(z\)’로부터 필요한 “remainder” 정보를 모델링하는 것을 목표로 하며, ‘\(I(x; z, z_n)\)’를 최대화함. 동시에 ‘\(z_n\)’은 ‘\(y\)’와 상관된 모든 정보를 제거하도록, 즉 ‘\(I(z_n; y)\)’를 최소화함. 따라서 ‘\(y\)’와 상관된 모든 정보는 complementary하게 ‘\(z\)’에 인코딩되어야 함. 이때 (2)에서의 capacity constraint 역할은 더욱 중요해짐. 단순히 ‘\(z\)’를 regularize하는 것뿐 아니라, ‘\(z_n\)’이 ‘\(y\)’ 예측에 불필요한 정보를 ‘\(z\)’로 밀어 넣는 것을 억제하여, (3)과 유사하게 ‘\(z\)’와 ‘\(z_n\)’ 간 경쟁 구도를 형성함. 기존 IB objective (2)와 결합하여, nuisance-extended IB(NIB)를 다음과 같이 정의함:
‘\(\max_f R_{NIB}(f) := R_{IB}(f) - I(z_n; y) + \gamma I(x; z, z_n)\)’ (4)
여기서 ‘\(\alpha \ge 0\)’임. 제안하는 NIB objective는 nuisance ‘\(z_n\)’을 도입하여 IB를 regularized한 형태로 볼 수 있음.
구체적으로, (4)는 추가적으로 ‘\(I(x; z, z_n)\)’을 최대화하고 ‘\(I(z_n; y)\)’를 최소화하도록 강제함. 즉, 각각 ‘\(H(x|z, z_n) = 0\)’ 및 ‘\(I(z_n; y) = 0\)’을 만족하도록 유도함.
다음 Lemma는 이러한 조건과 더불어 ‘\(z \perp z_n\)’ independence가 주어졌을 때, ‘\(I(\hat{z}; y)\)’로부터 원래의 ‘\(I(x; y)\)’를 복원할 수 있음을 보임(유도는 Appendix E 참조):
Lemma 1. random variable ‘\(x \in \mathcal{X}\)’, ‘\(y \in \mathcal{Y}\)’가 주어지고, ‘\(\hat{x}\)’가 ‘\(I(x; y) = I(\hat{x}; y)\)’를 만족하는 noisy observation이라고 하자. ‘\([\hat{z}, \hat{z}_n] := f(\hat{x})\)’가 (a) ‘\(H(\hat{x}|\hat{z}, \hat{z}_n) = 0\)’, (b) ‘\(I(\hat{z}_n; y) = 0\)’, (c) ‘\(\hat{z} \perp \hat{z}_n\)’를 만족하면, ‘\(I(\hat{z}; y) = I(x; y)\)’가 성립함.
이후 섹션에서는 autoencoder 기반 아키텍처를 통해 제안한 NIB의 practical design을 설명함. Section 3.1과 3.2에서는 각각 loss와 architecture를 상세히 다루고, Section 3.3에서는 전체 training 절차를 요약함. Figure 2는 전체 framework의 개요를 보여줌.
fig2 설명:
(1) Encoder (왼쪽 위 주황 블록)
- 입력 이미지 xx가 encoder EE로 들어감
- 출력이 두 갈래로 나뉨:
- zz: 의미/타깃 yy와 관련된 “semantic” representation
- znz_n: yy와 무관해야 하는 “nuisance” representation
(2) Decoder (왼쪽 아래 파랑 블록)
- (z,zn)(z, z_n)을 decoder GG에 넣어서 x^=G(z,zn)\hat{x}=G(z,z_n)를 복원함
- 그림의 초록색 “max log p(x|z, z_n)”가 이 부분임 → reconstruction loss가 걸림
여기서 reconstruction을 넣는 이유는 “생성모델 만들기”라기보다,
- 입력 xx의 정보가 z나 zₙ 중 어딘가엔 남아있게 강제
- 그 상태에서 znz_n이 yy 정보를 못 가지게 하면, yy 관련 정보가 zz로 몰리게 됨
이게 NIB의 핵심 목적임.
(3) 파란 MLP: z→yz \to y
- 오른쪽 파란 MLP가 zz로부터 yy를 예측함
- 옆의 파란 글씨 “max I(z;y) − βI(x;z)”가 바로 IB/VIB 쪽 목표를 요약한 것임
즉, zz는 잘 분류해야 하고(I(z;y)↑), 동시에 너무 많은 입력 정보를 그대로 들고 있지 않게(I(x;z)↓) 압축되는 방향임.
(4) 빨간 MLP: zn→yz_n \to y (하지만 “못 맞추게” 만듦)
- 오른쪽 빨간 MLP는 znz_n에서 yy를 예측하려고 시도하는 auxiliary classifier임
- 그런데 학습 목표는 “잘 맞추기”가 아니라, znz_n이 yy를 예측 못 하게 만드는 것임
- 그림의 빨간 글씨 “min I(z_n;y)”가 이 목적임
실제로 논문은 znz_n로부터의 예측이 uniform distribution에 가깝게 되도록 CE를 걸어버려서, znz_n에 label 정보가 남지 않게 밀어버림.
(5) 가운데 D_z: “real vs fake” (adversarial / GAN 스타일)
- DzD_z는 (z,zn)(z,z_n) 쌍을 보고
- encoder에서 나온 (z,zn)(z,z_n) = “real”
- N(0,I)\mathcal{N}(0,I)에서 샘플링한 (z,zn)(z,z_n) = “fake”
를 구분하려는 discriminator임
- 이건 “분포를 맞추는” 전형적인 adversarial(=GAN) 게임이라 GAN 방식이라고 부름
여기서 의도는:
- (z,zn)(z,z_n)의 joint prior가 isotropic Gaussian처럼 보이게 만들면
- 결과적으로 zz와 znz_n이 서로 섞여(의존/상관) 있는 걸 줄이고
- “의미 정보 vs nuisance 정보” 분리가 더 안정적으로 되게 하는 것임
논문은 IB를 이렇게 확장함:
- z = y 관련 정보
- zₙ = nuisance 정보
그리고 조건을 추가함:
- I(z_n; y) = 0 (nuisance는 y와 독립)
- I(x; z, z_n) 크게 (둘이 합쳐 x 복원 가능)
원래 IB(Information Bottleneck)는 이렇게 생겼다:
\max I(z; y) - \beta I(x; z)\
문제는:
- mutual information I(⋅;⋅)I(\cdot;\cdot)을 직접 계산하기가 거의 불가능함
- 특히 딥러닝에서는 고차원 확률분포라 계산이 어려움
그래서 등장한 것이 Variational Information Bottleneck (VIB)이다.
VIB의 실제 수식 형태
VIB에서 사용하는 loss는 보통 이렇게 생긴다:
L_{VIB} = \mathbb{E}_{\epsilon}[-\log q(y|f(x,\epsilon))] + \beta \, KL(p(z|x) \| r(z))
해석
첫 번째 항:
- 분류 loss (예측 정확도)
두 번째 항:
- KL divergence
- z가 너무 많은 정보를 담지 못하게 제한
- 입력 압축 역할
즉,
IB의 I(z;y)I(z;y)와 I(x;z)I(x;z)를 각각
cross entropy와 KL 항으로 근사한 것
IB를 딥러닝에서 실제로 학습 가능하도록 만든 variational 근사 버전
3.1. AENIB: A practical autoencoder-based design
(4)에서 정의한 NIB objective와 Lemma 1에 기반해, 제안한 framework를 구현하기 위한 practical training objective를 설계함.
여기서는 variational information bottleneck(VIB) [1] 위에 autoencoder 기반 아키텍처를 얹은 NIB의 단순한 instantiation을 제시하며, 이를 autoencoder-based NIB(AENIB)라 부름.
전반적으로 Lemma 1은 robust encoder ‘\(f\)’가 ‘\(\hat{z}\)’에 대해 세 가지 측면에서 generalization을 달성하는 “좋은” nuisance model을 요구한다고 말함.
세 측면은 (a) 좋은 reconstruction, (b) nuisanceness, (c) ‘\(z\)’와 ‘\(z_n\)’ 사이의 independence임.
이러한 behavior를 모델링하기 위해, encoder ‘\(f : \mathcal{X} \to \mathcal{Z}\)’뿐 아니라 decoder ‘\(g : \mathcal{Z} \to \mathcal{X}\)’도 함께 두고, autoencoder 기반 loss와 두 개의 adversarial loss [28]를 포함하는 다음 practical training objective를 채택함.
(a) 먼저 reconstruction loss를 두어 ‘\(\log p(x|z, z_n)\)’를 최대화함. 표준 설계에서는 decoder output이 ‘\(\mathcal{N}(x, \sigma^2 I)\)’를 따른다고 가정하며, 이는 mean-squared error(MSE)와 동치임. 여기서는 다른 loss들과의 균형을 효율적으로 맞추기 위해 normalized MSE를 사용함:2
‘\(L_{\text{recon}} := \frac{1}{\|x\|_2^2}\|x - g(z, z_n)\|_2^2\)’ (5)
역할:
(z, zₙ) 둘을 합치면 x를 복구할 수 있어야 한다
왜 필요?
- x의 정보가 다 z 또는 zₙ 중 하나에는 들어가 있어야 함
- 정보가 사라지지 못하게 만드는 장치
즉, 정보 분해를 강제함.
(b) ‘\(y\)’에 대한 ‘\(z_n\)’의 nuisance-ness를 강제하기 위해, ‘\(p(y|z_n)\)’을 multi-layer perceptron(MLP) ‘\(q_n\)’으로 근사하고 adversarial training을 수행함:
‘\(L_{\text{nuis}} := \mathbb{E}_x[\mathrm{CE}(q_n^{*}(z_n), \frac{1}{|\mathcal{Y}|})]\)’,
여기서 ‘\(q_n^{*} := \min_{q_n} \mathbb{E}_{x,y}[\mathrm{CE}(q_n(z_n), y)]\)’ (6)
이며, CE는 cross entropy loss를 의미함. 여기서는 ‘\(\mathcal{Y}\)’에서 uniform distribution을 향하도록 CE를 최적화함.3
역할:
zₙ으로는 y를 예측하지 못하게 하라
방법:
- zₙ → y 분류기를 붙임
- 그런데 uniform distribution 쪽으로 밀어버림
즉,
I(zn;y)≈0I(z_n; y) \approx 0
강제함.
→ nuisance에는 의미 정보 못 들어가게 함
(c) ‘\(z\)’와 ‘\(z_n\)’ 사이의 independence를 유도하기 위해, ‘\(z\)’와 ‘\(z_n\)’의 joint prior가 isotropic Gaussian,
즉 ‘\(p(z, z_n) \sim \mathcal{N}(0, I)\)’라고 가정하고, 2-layer MLP discriminator ‘\(q_z\)’를 사용한 GAN training을 수행함:
‘\(L_{\text{ind}} := \max_{q_z} \mathbb{E}_x[\log(q_z(f(x)))] + \mathbb{E}_{z,z_n \sim \mathcal{N}(0,I)}[\log(1 - q_z(z, z_n))]\)’ (7)
역할:
z와 zₙ이 서로 정보 공유하지 않게 하라
방법:
- GAN 형태로
- joint 분포가 Gaussian처럼 보이게 만듦
의미:
- z가 zₙ에 숨겨진 정보 빼먹지 못하게 함
- 둘이 역할 분담하게 만듦
왜 nuisance에서 uniform CE를 쓰냐?
AENIB에서 목표는:
z_n이 y를 예측하지 못하게 하는 것
즉,
- z_n로 만든 classifier 출력이
- 어떤 클래스도 특별히 높지 않게 만들고 싶음
그래서:
- one-hot target 대신
- uniform target을 줌
→ 모델이 모든 클래스에 대해 같은 확률을 내게 됨
즉, I(z_n; y) 강제하는 것임.
마지막으로, NIB (4)에서의 original IB objective ‘\(R_{IB}(f)\)’를 근사하기 위해, ‘\(R_{IB}\)’의 lower bound를 제공할 수 있는 variational information bottleneck(VIB) [1] objective ‘\(L_{\text{VIB}}^{\beta}\)’를 대신 최대화함.4 구체적으로 VIB는 (a) ‘\(p(y|z)\)’를 (parametrized) decoder neural network ‘\(q(y|z)\)’로 variational approximation하고, (b) ‘\(p(z)\)’를 isotropic Gaussian ‘\(\mathcal{N}(z|0,I)\)’ 같은 “더 쉬운” 분포 ‘\(r(z)\)’로 근사함. ‘\(f(x,\epsilon)\)’에 대해 Gaussian decoder (1)을 가정하면 다음을 얻음:
‘\(L_{\text{VIB}}^{\beta} := \frac{1}{n}\sum_{i=1}^{n}\mathbb{E}_{\epsilon}[-\log q(y_i|f(x_i,\epsilon))] + \beta \,\mathrm{KL}(p(z|x_i)\|r(z))\)’ (8)
2 Appendix C에서 SSIM 기반 [118] reconstruction loss도 탐색했으며, Transformer-based model에서 특히 robustness에 도움이 됨을 확인함.
3 대안으로 ‘\(\mathrm{CE}(q_n^{*}(z_n), y)\)’를 직접 최대화할 수도 있으나, 여기서는 maximization 기반 loss의 잠재적 instability를 피하기 위해 현재 설계를 사용함.
4 VIB framework(및 GAN)에 대한 더 자세한 설명은 Appendix F.2에 있음.
전체적 loss
- z는 y를 잘 맞춰야 함 (VIB)
- z는 너무 많은 정보를 담지 말아야 함 (VIB KL)
- (z, zₙ)은 x를 복구해야 함 (recon)
- zₙ은 y 정보를 가지면 안 됨 (nuis)
- z와 zₙ은 서로 섞이면 안 됨 (ind)
CE (Cross Entropy)란?
예측 분포와 타겟 분포 사이의 차이를 측정하는 loss
수식은 이렇게 생김:
CE(p,q)=−∑ipilogqiCE(p, q) = - \sum_i p_i \log q_i
여기서:
- = target 분포
- q = 모델이 예측한 분포
CNN의 마지막 feature만 쓰지 않고,
중간 레이어들의 평균/분산 통계를 모아서
z와 zₙ을 만들겠다
3.2. Architectures for nuisance modeling
원칙적으로 본 framework는 어떤 encoder architecture와도 일반적으로 호환됨.
예를 들어 encoder ‘\(f : \mathcal{X} \to \mathcal{Z}\)’와 decoder ‘\(g : \mathcal{Z} \to \mathcal{X}\)’를 둘 수 있음.
VIB를 적용하기 위해, encoder가 dimension ‘\(2K\)’의 두 output head를 갖는다고 가정하며, 여기서 ‘\(K\)’는 ‘\(z\)’의 dimension을 의미함.
각 output head는 reparametrization을 통해 Gaussian random variable을 모델링함.
즉, ‘\(z \in \mathbb{R}^{K}\)’와 ‘\(z_n \in \mathbb{R}^{K_n}\)’ 모두에 대해 (encoder output으로) ‘\((\mu,\sigma)\)’를 모델링함.
Gaussian을 표현하려면 뭐가 필요하나?
K차원 z를 표현하려면:
- 평균 벡터 μ∈RK\mu \in \mathbb{R}^K
- 분산 벡터 σ∈RK\sigma \in \mathbb{R}^K
두 개가 필요함.
즉, 총 2K개의 값이 필요함.
convention을 따라 deep feed-forward representation을 단순히 사용해 ‘\(f\)’가 ‘\(z\)’와 ‘\(z_n\)’을 모델링하는 것도 가능하지만, 표준 (discriminative) architecture에서 본질적으로 “generative”한 nuisance ‘\(z_n\)’을 모델링하면 training instability가 발생해 성능이 흔들리는 경향을 관찰함. nuisance information은 주어진 input의 더 fine detail을 모델링해야 하는 경우가 많으며, 이런 detail은 classification을 위한 later layer보다는 ‘\(f\)’의 early layer에 더 존재하는 경우가 많기 때문임.
본 논문에서는 convolutional network와 Vision Transformer(ViT) [24] 모두에서 nuisance modeling의 안정성을 높이기 위한 단순한 architectural treatment를 제안함.
본 section은 convolutional network 설계를 소개하며, patch-level representation이 제공되어 더 단순한 ViT 기반 설계는 Appendix C를 참조하라 함.
convolutional encoder ‘\(f\)’가 주어졌을 때, ‘\(f\)’의 최종 output을 직접 사용하기보다는 internal feature statistics의 collection으로부터 ‘\(z_n\)’(및 ‘\(z\)’)을 인코딩함.
구체적으로, 입력 ‘\(x\)’에 대해 ‘\(f(x)\)’로부터 ‘\(L\)’개의 intermediate feature map ‘\(x^{(1)}, \cdots, x^{(L)}\)’를 추출하고, 다음과 같이 ‘\(x\)’의 projection을 정의함:
‘\(\Pi_f(x) := \begin{bmatrix} m^{(1)} & m^{(2)} & \cdots & m^{(L)} \\ s^{(1)} & s^{(2)} & \cdots & s^{(L)} \end{bmatrix}\)’ (9)
여기서 ‘\(m^{(l)}\)’과 ‘\(s^{(l)}\)’는 ‘\(x^{(l)}\)’의 feature map에 대한 1차/2차 moment이며, ‘\(x^{(l)} \in \mathbb{R}^{HWC}\)’라고 가정할 때
‘\(m_c^{(l)} := \frac{1}{HW}\sum_{hw} x_{hwc}^{(l)}\)’, ‘\(s_c^{(l)} := \frac{1}{HW}\sum_{hw}(x_{hwc}^{(l)} - m_c^{(l)})^2\)’
로 정의됨.
실제로는 이걸 한 줄로 flatten해서 큰 벡터로 씀.
핵심: 초기~중기 레이어의 low-level 통계가 nuisance(스타일/텍스처/배경) 정보를 잘 담는다는 경험적 이유로 이렇게 함.
projection(사영) Πf(x)는 수학적 의미의 “직교 사영”이 아니라, 중간 feature들을 하나의 벡터로 요약해서 쓰기 위한 정리 작업
CNN encoder의 중간 feature map들은 보통 이런 형태임:
x^{(l)} \in \mathbb{R}^{H_l \times W_l \times C_l}
- 공간 차원(H, W)
- 채널 차원(C)
을 가진 3차원 텐서임.
문제는: 이런 텐서를 그대로 MLP에 넣기 어렵고,레이어마다 크기도 다름.
그래서
- 공간 차원(HW)을 없애고
- 각 레이어를 하나의 고정 길이 벡터로 만들 필요가 있음.
이 역할이 projection임.
왜 모든 통계를 모으고 그걸 MLP에 넣냐?
CNN 버전에서 Πf(x)는:
- 중간 feature들의 평균/분산 통계 벡터
이걸 그대로 z로 쓰지 않고 MLP에 넣는 이유는:
- 차원 축소
- 비선형 변환
- z와 zₙ을 서로 다른 역할로 만들기
통계 벡터 → MLP → z용 분포
통계 벡터 → 다른 MLP → zₙ용 분포
이렇게 “역할 분리”를 해주는 것임.
Appendix I에서, 이 simple projection이 ‘\(x\)’의 generative representation을 충분히 인코딩할 수 있음을 보임.
즉, ‘\(\Pi_f\)’ 위에 정의된 discriminator로 GAN을 성공적이고 효율적으로 학습할 수 있음을 보임.
이 관찰에 기반해, encoder representation ‘\(z\)’와 ‘\(z_n\)’을 모델링할 때 ‘\(\Pi_f\)’를 채택함.
(9)의 ‘\(\Pi_f(x)\)’에 MLP를 단순히 적용하는 방식으로 ‘\(z\)’와 ‘\(z_n\)’을 인코딩함.
단순함에도 불구하고, 이 treatment가 AENIB의 stable training을 가능하게 함을 관찰함.
3.3. Overall training objective
제안한 objective들과 VIB loss ‘\(L_{\text{VIB}}^{\beta}\)’ (8)을 결합하면 최종 objective로 이어짐.
실무적으로 multiple loss를 결합하면 additional hyperparameter가 생길 수 있으나, reconstruction loss ‘\(L_{\text{recon}}\)’과 original VIB loss의 ‘\(\beta\)’를 제외하면 대부분의 제안 loss는 scaling 없이 더해도 됨을 확인함.
따라서 다음을 얻음:
‘\(L_{\text{AENIB}} := L_{\text{VIB}}^{\beta} + \gamma \cdot L_{\text{recon}} + L_{\text{nuis}} + L_{\text{ind}}\)’ (10)
Appendix A의 Algorithm 1이 procedure를 요약함.
4. Experiments
제안한 AENIB training의 효과를 model reliability의 다양한 측면에서 검증함. 구체적으로 (a) corruption 및 natural robustness(Section 4.1), (b) novelty detection(Section 4.2), (c) certified adversarial robustness(Section 4.3) task를 다루며, 이들 모두는 task-specific prior 없이 [37, 39, 82] 특히 challenging했던 설정임. component-wise analysis를 위한 ablation study는 Appendix D에 제시함. 또한 Appendix I에서는 generative modeling 맥락에서 제안 component들의 평가도 제시함. dataset, training detail, hyperparameter 등 experiment의 full detail은 Appendix B에 있음.
4.1. Robustness against natural corruptions
먼저 본 방법의 corruption robustness, 즉 natural corruption(fog, brightness 등) 및 인간에게 여전히 semantic하게 보이는 distribution shift 하에서의 generalization capability를 평가함. 이를 위해 generalization을 측정하기 위한 목적으로 CIFAR-10과 ImageNet에서 파생된 광범위한 benchmark를 고려함. 즉, CIFAR-10 model에 대해서는 (a) CIFAR-10/100-C [35]를 테스트함. 이는 15개의 common corruption을 5개의 severity level로 시뮬레이션한 CIFAR-10/100의 corrupted version임. 또한 (b) CIFAR-10.1 [94], CIFAR-10.2 [81], CINIC-10 [21]도 테스트함. 이는 CIFAR-10 test set을 세 번 재생성한(regeneration) 데이터셋임. 반면 ImageNet model에 대해서는 (a) ImageNet validation set의 corrupted version인 ImageNet-C [35], (b) 200개 ImageNet class에 대한 rendition image collection인 ImageNet-R [34]과 ImageNet-Sketch [117], 그리고 (c) background change에 대한 model bias를 평가하기 위한 Background Challenge [119] benchmark를 테스트함. 본 section은 주로 ViT [24, 111] 기반 architecture 결과를 보고하지만, ResNet-18 [33] 결과도 Appendix H에서 보고함.
Table 1과 2는 CIFAR 기반 model 결과를 요약함. Table 1에서 AENIB가 VIB 대비 corruption error를 유의미하고 일관되게 개선함을 관찰함. 또한 이러한 gain은 state-of-the-art 방법들과 비교해도 강력함. 예를 들어 AENIB만으로도 강력한 baseline인 AugMix [39]를 outperform 가능함. 더 최근 방법인 PixMix [40]는 extra (pattern-like) data를 활용해 더 낮은 corruption error를 달성할 수 있지만, (a) “AENIB + PixMix”에서 보이듯 AENIB도 PixMix(즉 extra data)로부터 이득을 얻으며, (b) Table 2 결과는 AENIB의 generalization capability가 CIFAR-10.1, 10.2, CINIC-10에서 PixMix보다 더 좋음을 보여줌. 즉, domain-specific data에 덜 의존하면서 common corruption을 넘어선 영역에서도 더 잘 generalize함.
다음으로 Table 3는 AENIB의 효과가 ImageNet이라는 더 large-scale, higher-resolution dataset으로 generalize됨을 강조함. 추가적인 data augmentation을 training 중에 더 활용하지 않아도, AENIB가 다양한 corruption type에 대해 robust accuracy를 일관되게 개선함을 여전히 관찰함. Figure 3는 다양한 data augmentation과 hyperparameter 전반에서 Cross-entropy와 AENIB가 만들어내는 linear trend를 비교하며, 서로 다른 model들에서 in-distribution vs out-of-distribution 성능 간 상관관계에 대한 최근 관찰 [34, 87, 105]을 고려했을 때, AENIB가 effective robustness [105] 관점에서도 더 나은 operating point를 보임을 확인해 줌.
마지막으로 Table 4는 Background Challenge [119]에서 ImageNet classifier를 추가로 평가함. 이 benchmark는 background shift에 대한 model robustness를 테스트하기 위해 만들어졌으며, 구체적으로 ImageNet-9(ORIGINAL, ImageNet의 370개 subclass를 결합함)의 variant를 서로 다른 background 조합으로 구성함. AENIB 기반 model은 이 benchmark에서도 cross-entropy baseline 대비 일관되게 성능을 개선하며, AENIB가 background change에 덜 biased된 feature를 학습하는 경향이 있음을 보여줌.
4.2. Novelty detection
다음으로, AENIB model이 out-of-distribution sample(OOD)을 잘 탐지하는 detector로도 동작할 수 있음을 보이며, 즉 novelty detection task를 해결할 수 있음을 보임. 일반적으로 이 task는 novel sample과 in-distribution sample을 구분하려는 binary classification problem으로 정의됨. 여기서의 전형적인 practice는 각 input에 대해 score function을 정의하고(예: maximum confidence score [36]), score가 낮을 때 해당 sample을 out-of-distribution으로 thresholding하는 방식임. AENIB model을 위한 score function을 정의하기 위해, 먼저 standard model에는 없고 AENIB에서만 가능한 ‘\(z_n\)’의 log-likelihood가 in-distribution에서 semantic하게 멀리 떨어진 novelty를 탐지하는 강력한 score가 될 수 있음을 관찰함. 구체적으로 ‘\(z\)’가 isotropic Gaussian ‘\(\mathcal{N}(0, I)\)’를 따른다고 가정하므로, ‘\(\log \mathcal{N}(z_n; 0, I) = -\frac{1}{2}\|z_n\|^2\)’를 사용함. 소위 “harder” novelty를 탐지하기 위해, 단순히 ‘\(\max_y p(y|x)\)’를 쓰는 대신, parameter ‘\(\alpha > 0\)’를 갖는 symmetric Dirichlet distribution 하에서의 ‘\(y\)’ log-likelihood score를 사용할 것을 제안함. 즉, ‘\(\mathrm{Dir}_\alpha(y) \in \Delta^{|\mathcal{Y}|-1}\)’에 대해 ‘\(\log \mathrm{Dir}_\alpha(y) = (\alpha - 1)\sum_i \log y_i\)’를 사용함. ‘\(\alpha \to 0\)’일수록 이 분포는 symmetric (discrete) one-hot distribution에 가까워지며, 이는 대부분 classification task에서 타당함. 여기서는 실험 전반에서 단순히 ‘\(\alpha = 0.05\)’를 사용함.5
평가 benchmark로 두 가지를 고려함. (a) “standard” benchmark는 기존 문헌 [36, 72, 74]에서 활발히 채택되어 왔으며, CIFAR-10 test set을 in-distribution으로 두고 다른 독립 dataset들에 대한 detection performance를 측정함. (b) 반면 최근 OBJECTS benchmark [122]는 CIFAR-10 benchmark를 확장하여 OOD 평가에서 “near” in-distribution까지 고려함. 구체적으로 OBJECTS는 test-time에서 CIFAR-10뿐 아니라 CIFAR-10-C [35]와 ImageNet-10도 in-distribution으로 가정하여 detection을 훨씬 더 challenging하게 만듦. 본 실험에서는 [122]의 setup을 따라 CIFAR-10에서 학습한 ResNet-18 [33] model들을 비교함.
결과는 standard benchmark와 OBJECTS benchmark에 대해 각각 Table 5와 Table 6에 보고됨. 전반적으로, AENIB의 ‘\(z_n\)’과 ‘\(y\)’ 정보를 결합한 score function이 강력한 baseline 대비 novelty detection을 complementary하게 유의미하게 개선함을 확인함. 이는 nuisance modeling의 효과를 보여줌. 예를 들어 Table 5에서, combined score는 SVHN, LSUN, ImageNet dataset을 탐지하는 데 거의 perfect에 가까운 AUROC을 달성함. 반면 Table 6에서는, AENIB가 OBJECTS vs. MNIST 설정에서 이전 최고 AUROC(Mahalanobis [72])을 77.04 → 92.43으로 개선함. OBJECTS에서의 개선은 AENIB에서 얻은 representation과 score 모두가 더 semantic한 관점에서 in-distribution과 out-of-distribution을 더 잘 구분하도록 도움을 준다는 것을 시사함.
4.3. Certified adversarial robustness
또한 adversarial robustness [29, 82, 102]를 평가하며, 주어진 representation에 대한 certified robustness를 측정할 수 있는 randomized smoothing framework [17, 70]를 채택함. 구체적으로 어떤 classifier든 Gaussian noise 하에서 prediction을 averaging함으로써 robustified될 수 있으며, 입력 ‘\(x\)’에서의 robustness는 classifier가 ‘\(\mathcal{N}(x, \sigma^2 I)\)’를 얼마나 일관되게 분류하는지에 달려 있음 [50]. 이 평가 protocol 하에서, randomized smoothing을 AENIB와 결합하면 adversarially-robust representation이 자연스러운 byproduct로 얻어질 수 있다고 제안함. 이는 종종 큰 training cost가 들고 robustness에 대한 certification이 없는 thorough adversarial training method [82]를 사용하지 않아도 된다는 의미임.
standard certification protocol [17]을 따라 radius ‘\(r\)’에서의 certified test accuracy를 비교함. 이는 smoothed classifier가 test sample을 정확히 분류하면서 certified radius가 ‘\(r\)’보다 큰 sample의 비율로 정의됨. CIFAR-10에서 ViT-S model을 고려하고, 본 실험에서는 ‘\(\sigma = 0.1\)’로 가정함. Figure 4에 요약된 결과는 AENIB가 테스트한 모든 radius에서 baseline 대비 유의미하게 더 좋은 certified robustness를 달성함을 보여줌. 예를 들어 ‘\(\epsilon = 0.1\)’에서 VIB의 certified robust accuracy를 39.6% → 56.8%로 개선함. 다시 말해, AENIB에서 얻은 robustness는 threat model에 대한 specific knowledge에서 나온 것이 아니며, 이는 randomized smoothing과 결합할 때 AENIB가 “공짜” adversarial robustness를 제공할 수 있음을 시사함. 이는 AENIB의 robustness가 단지 크기만 큰 것이 아니라 input별로도 일관적임을 확인해 주며, 특히 더 큰 ‘\(r\)’에서 높은 certified robustness를 보인다는 점을 고려하면 더 의미가 큼.
4.4. Comparison with DisenIB [92]
본 section에서는 AENIB를 관련 연구인 DisenIB [92] 및 VIB의 다른 variant들인 Nonlinear-VIB [65], Squared-VIB [106]와 비교함. 여기서 DisenIB는 supervised disentangling을 목적으로 nuisance modeling(FactorVAE [57] 기반)도 함께 고려하는 IB variant임. 구체적으로 DisenIB는 서로 independent한 두 encoder ‘\(z := f(x)\)’와 ‘\(z_n := g(x)\)’를 두고, 다음 objective를 최적화함:
‘\(\max_{f,g}\; I(z; y) + I(x; z_n, y) - I(z_n; z)\)’ (11)
제안한 NIB (4)와 비교할 때, 두 objective의 가장 중요한 차이는 “reconstruction” term에 있음. 즉, (11)의 ‘\(I(x; z_n, y)\)’ vs. 우리 (4)의 ‘\(I(x; z, z_n)\)’임. 이 차이 때문에 DisenIB objective (11)는 optimum에서도 ‘\(z\)’가 ‘\(x\)’의 “shortcut” signal(‘\(y\)’와 상관된)을 소수만 encode하는 경우를 배제하지 못함. 이는 ‘\(z\)’가 ‘\(x\)’의 가능한 모든 ‘\(y\)’-correlated signal을 최대한 encode하도록 만들려는 NIB (4)의 핵심 motivation과 대조됨.
실험 비교는 (a) 우리의 CIFAR-10 setup(Table 2)과 (b) MNIST [69]에서 DisenIB의 official implementation6을 직접 사용한 비교(Appendix G의 Table 9)로 수행함. 여기서는 benchmark를 확장하여 MNIST의 corrupted version인 MNIST-C [88]도 포함함. 후자의 비교에서는 encoder architecture로 단순한 4-layer convolutional network를 사용함. MNIST model은 모두 100K update 동안 학습하고, 나머지 training detail은 CIFAR 실험과 동일하게 따름(Appendix B.1 참조).
전반적으로 AENIB의 효과가 이러한 benchmark들에서도 여전히 적용됨을 관찰함. 이는 DisenIB와 대조적임. DisenIB의 효과는 예를 들어 [92]가 수행한 Gaussian noise를 OOD로 탐지하는 AUROC gain처럼 특정 설정에서는 보일 수 있으나, CIFAR-10-C나 MNIST-C로는 더 이상 일반화되지 못했기 때문임. 반면 AENIB는 이러한 corrupted benchmark에서도 여전히 개선을 보이며, 동일한 OOD task에서 perfect score도 달성함. Figure 5a에서는 DisenIB가 nuisance ‘\(z_n\)’에 highly semantic한 정보를 자주 남기는 것을 정성적으로도 관찰함. 즉, reconstruction이 다른 sample의 ‘\(z_n\)’과 swap하는 것만으로도 완전히 바뀔 수 있음. 이는 Figure 5b에서 비교하듯, AENIB가 본질적으로 해결하려는 문제임.
5 실험에서는 전반적으로 ‘\(\alpha = 0.05\)’를 사용함.
6 DisenIB의 official implementation을 의미함.
5. Related work
Out-of-distribution robustness. out-of-distribution input에 대한 neural network의 fragility를 드러낸 seminal work [2, 90, 102] 이후, robustness의 다양한 개념을 식별하고 개선하려는 상당한 시도가 이루어져 왔음. 예를 들어 novel input 탐지 [36, 71, 72, 103], corruption에 대한 robustness [27, 35, 39, 119], adversarial noise에 대한 robustness [5, 11, 17, 82] 등이 있음. 그러나 neural network가 extrapolate하도록 만드는 근본적인 어려움 때문에, robustness 문헌에서의 대부분의 발전은 개별 문제와 밀접하게 관련된 prior를 가정하는 조건 하에서 이루어져 왔음. 예를 들어 external data나 data augmentation [37, 39], test-time sample로부터의 extra information [116], 혹은 threat model에 대한 specific knowledge [52, 112] 등을 가정함. 본 연구에서는 이러한 prior를 가정하지 않고, test-time의 noisy observation 하에서 standard information bottleneck 원리를 확장한 새로운 training scheme를 통해 multiple robustness notion을 동시에 개선하고자 함.
Hybrid generative-discriminative modeling. 제안한 방법은 generative model을 결합하여 discriminative model의 robustness를 개선하는 새로운 접근으로도 볼 수 있음. 이는 최근 연구 [31, 72, 99, 123]에서 탐구된 맥락과 연결됨. 예를 들어 [72, 73]은 simple하지만 generation에 대해서는 low expressivity를 갖는 Gaussian mixture model을 discriminative classifier에 결합함. Joint Energy-based Model(JEM) [31, 123] 계열 연구는 energy-based model을 가정하지만, 그 목적을 위해 notable한 training instability를 겪음. 본 연구에서는 이러한 instability를 피하기 위해 autoencoder 기반 model을 제안하고, nuisance가 주어진 discriminative representation을 succinct하게 보완하여 generative하게 만들도록 설계함. 우리의 접근은 두 세계의 장점을 동시에 취함을 보임. 즉, (a) stable training을 가능하게 하면서도, (b) 높은 expressive generative performance를 달성함.
Nuisance modeling. nuisance를 도입하는 아이디어는 invertible modeling, 즉 flow-based model [6, 23, 30, 59] 맥락에서도 고려될 수 있음. 여기서 nuisance는 관심 subspace에 대해 (full-information) encoding ‘\(z\)’를 분할함으로써 정의될 수 있음 [4, 46]. 이러한 접근과 달리, 우리의 autoencoder 기반 nuisance modeling은 arbitrary input에 대한 “full” invertibility에 초점을 두지 않고, 주어진 data manifold를 invert하는 데 초점을 둠. 이는 (a) 실제에서 훨씬 flexible한 encoder design을 가능하게 하고, (b) MNIST-scale [46]을 넘어서는 등 보다 scalable한 nuisance ‘\(z_n\)’의 generative modeling을 가능하게 함. 다른 관련 연구 [48, 49, 92] 역시 nuisance factor를 위한 encoder를 도입하나, nuisance-ness의 개념은 robustness 목적에서 ‘\(y\)’와의 independence가 아니라 feature disentangling 목적에서 ‘\(z\)’와의 independence에 초점을 둠. 예를 들어 DisenIB [92]는 semantic embedding과 nuisance embedding 사이의 independence를 강제하기 위해 FactorVAE [57]를 적용함. 그러나 이러한 아이디어가 MNIST를 넘어 확장 가능한지에 대해서는 문헌에서 의문이 제기되어 왔음. 본 연구는 최근 architecture와 dataset을 활용한 practical design을 탐구하고 정립함으로써, 다양한 modern security metric을 다루는 설정에서 이를 확장 가능함을 보임.
관련 연구에 대한 보다 광범위하고 상세한 논의는 Appendix F에 제공됨.
6. Conclusion
좋은 nuisance model을 갖는 것이 reliable representation을 유도하는 실질적인 접근이 될 수 있음을 제안함. 데이터로부터 deep nuisance representation을 학습하는 practical method를 개발하고, prior를 가정하지 않는 challenging한 setup [105] 하에서 다양한 reliability measure를 개선하는 데 효과적임을 보임. 본 연구가 deep learning에서 out-of-distribution generalization을 더 잘 이해하는 방향으로 나아가는 유의미한 단계가 될 수 있다고 믿음. 현재 scope는 autoencoder 기반 model의 특정 설계에 초점을 두고 있으나, nuisance-extended IB framework 자체는 이에 국한되지 않으며, 향후 연구에서는 보다 다양한 구현을 고려할 수 있음. 궁극적으로 우리는 mutual information constraint를 갖는 도전적인 형태의 adversarial training을 근사하고자 하며, 이는 향후 탐구할 유망한 방향이라고 믿음.