🤖 machine learning

Bayesian Tensor Decomposition with Diffusion Model Prior

이 논문은 사전 학습된 확산 모델을 데이터 사전으로 통합하고 자동 랭크 선택을 위한 누적 수축 과정을 활용하며, 트랙터블한 추론을 가능하게 하는 분할 깁스 샘플러를 사용하여 심각한 손상 상황에서의 이미지 인페인팅 및 디노이징에서 우수한 성능을 달나하는 베이지안 텐서 분해 프레임워크인 DiffBCP를 소개한다.

원저자: Zerui Tao, Qibin Zhao

게시일 2026-06-03

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Zerui Tao, Qibin Zhao

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신에게 도시나 얼굴 사진 같은 복잡한 이미지(텐서라고 불리는 3D 퍼즐)가 있다고 상상해 보세요. 불행하게도, 누군가가 이 퍼즐 조각의 거대한 덩어리들을 뜯어갔고, 남은 조각들은 노이즈(static과 scratches)로 뒤덮여 있습니다. 당신의 목표는 원래의 이미지가 어떤 모습이었는지 알아내는 것입니다.

이 논문은 이 퍼즐을 풀기 위한 새로운 방법인 DiffBCP를 소개합니다. 이 방법이 어떻게 작동하는지 쉬운 개념으로 나누어 설명하겠습니다.

1. 기존 방식: "저계수(Low-Rank)"를 통한 추측

전통적으로 컴퓨터는 이미지가 단순한 기저 구조를 가지고 있다고 가정하며 이 문제를 해결하려 했습니다. 그들은 이를 **"저계수성(Low-Rankness)"**이라고 부릅니다.

비유: 이미지가 몇 개의 뚜렷한 붓터치로 만들어진 그림이라고 상상해 보세요. 만약 그 몇 가지 붓터치의 패턴을 알고 있다면, 빠진 부분을 추측할 수 있습니다.
문제점: 이 방식은 이미지가 깨끗할 때는 매우 잘 작동합니다. 하지만 이미지가 심하게 손상되었거나 거대한 덩어리가 사라진 경우에는, 단순히 "단순하다"고 가정하는 것만으로는 부족합니다. 컴퓨터는 막히게 되고, 결과물은 흐릿하고 이상하게 나옵니다.

2. 새로운 재료: "확산 모델(Diffusion Model)" (예술 전문가)

저자들은 "저계수성"이 좋은 가이드라인이긴 하지만, 지저도 있고 복잡한 실제 데이터에는 충분하지 않다는 것을 깨달았습니다. 그래서 그들은 훨씬 더 똑똑한 조력자인 확산 모델을 추가하기로 했습니다.

비유: 확산 모델을 수백만 장의 사진을 본 세계적인 예술 전문가라고 생각해보세요. 이 전문가는 수학만 아는 것이 아니라, 현실적인 얼굴, 나무, 혹은 건물이 어떠해야 하는지에 대한 직관적인 "감각"을 가지고 있습니다.
혁신: 보통 이러한 예술 전문가들은 수학적 퍼즐 해결사와 결합하기가 어렵습니다. 저자들은 수학적 규칙을 깨뜨리지 않으면서도 이 전문가가 퍼즐 해결 과정을 안내할 수 있는 방법을 찾아냈습니다.

3. 하이브리드 팀: DiffBCP

새로운 방법인 DiffBCP는 "저계수" 수학과 "예술 전문가"의 팀업입니다.

협업 방식:
1. 수학 (저계수): 구조를 체계적으로 유지하고 퍼즐 조각들이 논리적으로 서로 맞물리도록 합니다. 또한, 실제로 필요한 "붓터치"가 얼마나 되는지 자동으로 결정하는 스마트한 필터 역할을 하여 너무 복잡해지지 않도록 합니다.
2. 전문가 (확산): 엉망이 되고 불완전한 퍼즐을 보고 이렇게 말합니다. "이 빠진 부분은 무작위한 덩어리가 아니라 창문처럼 보여야 해요." 그리고 그 빈 공간을 현실적인 질감으로 채워 넣습니다.

4. 핵심 비법: "분리 깁스 샘플러(Split Gibbs Sampler)"

이 둘을 섞는 것은 까다롭습니다. 왜냐하면 한쪽은 엄격한 수학을 말하고, 다른 한쪽은 확률을 말하기 때문입니다(서로 다른 언어를 사용합니다). 이들이 서로 소통하게 만들기 위해, 저자들은 **"분리 깁스 샘플러"**라는 특별한 통신 프로토콜을 발명했습니다.

비유: 두 사람이 집을 짓고 있다고 상상해 보세요. 한 명은 엄격한 건축가(수학)이고, 다른 한 명은 창의적인 인테리어 디자이너(확산 모델)입니다.
- 그들은 서로 다투는 대신, 번역가(샘플러)를 사용합니다.
- 건축가가 뼈대를 세웁니다.
- 번역가가 뼈대를 디자이너에게 전달하면, 디자이너가 현실적인 벽지와 가구를 추가합니다.
- 번역가가 이를 다시 건축가에게 전달하면, 건축가는 새로운 가구에 맞춰 뼈대를 조정합니다.
- 이 과정을 완벽한 집이 될 때까지 반복합니다.
이점: 이를 통해 컴퓨터는 노이즈를 자동으로 처리할 수 있습니다. 인간이 "노이즈 수준이 5%입니다"라고 알려줄 필요가 없습니다. 시스템이 작업하는 동안 스스로 노이즈 수준을 파악합니다.

5. 결과: 무엇을 발견했는가?

저자들은 다음 테스트를 수행했습니다:

이미지 인페인팅(Image Inpainting): 사진의 빠진 부분을 채우는 것 (예: 군중 속에서 사람을 제거하거나 찢어진 사진을 복구하는 것).
디노이징(Denoising): 입자가 거칠고 정적(static)이 가득한 이미지를 깨끗하게 만드는 것.
고해상도 및 분포 외 데이터(High-Resolution & Out-of-Distribution): 거대한 고해상도 이미지나, 전문가가 학습한 데이터와 전혀 다른 모습(예: 얼굴로 학습된 전문가에게 밤의 도시 풍경을 보여줌)을 가진 이미지에서도 이전 방식들보다 더 뛰어난 성능을 보였습니다.

요약하자면: 그들은 수학의 구조적 논리와 AI 예술 생성기의 창의적 직관을 결합한 시스템을 구축했습니다. 이를 통해 손상이 심하거나 이미지가 매우 큰 경우에도 이전 방법들보다 훨씬 더 정확하고 현실적으로 손상된 이미지를 재구성할 수 있게 되었습니다.

기술 요약: 확산 모델 사전 확률을 이용한 베이지안 텐서 분해

문제 정의

저계수 텐서 분해(Low-rank tensor decomposition, TD)는 데이터가 깨끗하고 완전히 관측되었을 때 다차원 데이터 분석에 효과적인 표준 도구이다. 그러나 심각한 결측치나 노이즈가 존재하는 상황에서는 성능이 크게 저하된다. 저계수 가정이 구조적 사전 확률(structural prior) 역할을 수행하지만, 실제 세계의 복잡한 데이터를 복구하기에는 불충분한 경우가 많다. 반대로, 최첨단 확산 모델(diffusion models)은 강력하고 학습된 데이터 사전 확률을 제공하지만, 현재의 TD 프레임워크와 쉽게 호환되지 않으며, 특히 저계수 제약 및 자동 계수(rank) 선택과 결합될 때 추론이 까다롭다는 문제가 있다.

방법론: DiffBCP

저자들은 학습된 확산 사전 확률과 구조적 저계수 제약을 통합하도록 설계된 하이브리드 사전 확률 기반 베이지안 CP 분해 프레임框架인 DiffBCP를 제안한다.

1. 확률적 정식화

본 모델은 관측된 데이터 $Y$ , 재구성된 텐서 $X$ , 그리고 잠재 CP 인자( $\lambda, A^{(1:N)}$ )에 대한 결합 확률 분포를 다음과 같이 정의한다:
$p(Y, X, \lambda, A^{(1:N)}, \tau) \propto p(Y | X, \tau) \cdot p(X | A^{(1:N)}, \lambda) \cdot p(A^{(1:N)}) \cdot p(\lambda) \cdot p(X) \cdot p(\tau)$

주요 구성 요소는 다음과 같다:

우도(Likelihood): 관측된 엔트리가 정밀도 $\tau$ 를 가진 정규 분포를 따른다고 가정한다.
구조적 제약: $p(X | A^{(1:N)}, \lambda)$ 는 델타 함수를 통해 CP 분해 구조( $X = \text{CP}(\lambda, A^{(1:N)})$ )를 강제한다.
잠재 인자 사전 확률(Latent Factor Priors):
- CUSP 사전 확률: CP 가중치 $\lambda$ 는 누적 수축 프로세스(Cumulative Shrinkage Process, CUSP) 사전 확률를 따른다. 이 계층적 사전 확률은 계수가 증가함에 따라 중복된 성분을 작은 상수 $\theta_\infty$ 로 수축시켜 자동 계수 선택을 가능하게 한다.
- 가우시안 사전 확률: 인자 행렬 $A^{(n)}$ 은 노름(norm)을 규제하기 위해 표준 가우시안 사전 확률를 따른다.
확산 사전 확률(Diffusion Prior): 재구성된 텐서 $X$ 에는 강력하고 사전 학습된 확산 모델 사전 확률 $p(X)$ 가 할당된다. 이는 저계수 가정만으로는 포착할 수 없는 복잡한 질감과 구조를 포착할 수 있는 암묵적이고 데이터 중심적인 규제 역할을 한다.

2. 추론: 분할 깁스 샘플러 (Split Gibbs Sampler, SGS)

우도, 저계수 제약, 그리고 암묵적인 확산 사전 확률 사이의 결합으로 인해 직접적인 사후 확률 추론은 불가능하다. 저자들은 이 항들을 분리하기 위해 보조 변수 $Z$ 를 사용하는 **분할 깁스 샘플러(Split Gibbs Sampler)**를 채택한다:

증강(Augmentation): 결합 항 $\phi(Z, X; \rho) = \frac{1}{2\rho^2}\|Z - X\|_F^2$ 를 갖는 보조 변수 $Z$ 를 도입한다.
잠재 변수 업데이트: $Z$ 가 주어졌을 때, CP 인자( $\lambda, A^{(n)}$ )와 하이퍼파라미터( $\tau, \theta, \nu$ )는 닫힌 형태의 공액 업데이트(Gaussian, Gamma, Beta 등)를 허용한다.
확산 블록 업데이트: $X$ 가 주어졌을 때, 보조 변수 $Z$ 는 확산 사전 확률과 결합 항에 의해 정의된 분포로부터 샘플링된다. 이는 저계수 추정치 $X$ 에 의해 유도되는 사전 학습된 스코어 네트워크 $s_\psi$ 를 통한 SDE(확률 미분 방정식) 역과정의 디노이징 문제로 해결된다.
노이즈 적응형 결합: 추론된 노이즈 정밀도 $\tau$ 에 따라 결합 파라미터 $\rho$ 가 동적으로 조정되는 스케줄( $\tau\rho^2 = c$ )이 도입된다. 이를 통해 수동 아닐링 스케줄에 대한 민감도를 줄인다.

3. 계수 적응 (Rank Adaptation)

본 프레임워크는 양방향 계수 적응 기법을 포함한다. CUSP 사전 확률은 비활성 성분을 수축시키고, 별도의 메커즘은 현재 계수가 불충분할 경우 새로운 성분을 추가한다. 이는 초기 계수 설정 오류에 대한 견고성을 보장한다.

주요 기여

하이브리드 사전 확률 프레임워크: 오프더쉘레프(off-the-shelf) 사전 학습된 확산 모델을 재구성된 텐서에 대한 데이터 사전 확률로 사용하면서, 동시에 구조적 저계수 CP 사전 확률과 결합한 최초의 베이지안 텐서 분해 모델이다.
다루기 쉬운 추론: 우도, 저계수 제약, 확산 사전 확률 사이의 복잡한 결합에도 불구하고 사후 추론을 가능하게 하는 분할 깁스 샘플러를 개발하였다. 여기에는 잠재 인자를 위한 닫힌 형태의 업데이트와 저계수 유도 디노이징을 포함한다.
자동 계수 및 노이즈 추론: CUSP 사전 확률과 완전 베이지안 노이즈 추정을 사용하여 계수와 노이즈 수준을 위한 수동 튜닝의 필요성을 제거하였다.
견고성: 적응형 결합 스케줄을 통해 초기 계수 설정 오류 및 다양한 노이즈 수준에 대한 견고성을 입증하였다.

실험 결과

본 방법은 FFHQ 및 ImageNet 데이터셋을 사용한 이미지 인페인팅(inpainting) 및 디노이징 작업, 그리고 고해상도 OOD(분포 외) 이미지에 대해 평가되었다.

성능: DiffBCP는 전통적인 베이지안 TD(BCP, BTR), 딥러닝 기반 TD(HLRTF, DeepTensor), 그리고 플러그 앤 플레이 방식(GLON, PnP-DM)을 포함한 베이스라인 모델들을 일관되게 능가했다.
- 90% 결측(Uniform mask) 조건의 FFHQ 데이터셋에서 DiffBCP는 28.28 dB PSNR을 달성하여, 가장 우수한 베이스라인인 DeepTensor보다 2.17 dB 향상된 성능을 보였다.
- 불규칙한 마스크(Irregular mask)의 경우, 성능 향상 폭은 더욱 두드러졌다 (DeepTensor 대비 +2.33 dB).
OOD 및 고해상도: 256×256 이미지로 학습된 확산 모델을 사용하여 2048×2048 해상도의 이미지(예: Marseille, Tokyo)를 성공적으로 재구성하였으며, 다른 방법들(예: PuTT)이 아티팩트를 생성하거나 실패하는 지점에서도 일관성과 세부 사항을 유지하였다.
절제 연구(Ablation Studies):
- 확산 사전 확률을 제거했을 때 구조적 마스크(Stripe, Irregular)에서 성능이 크게 저하되었으며, 이는 공간적 연속성을 회복하는 데 있어 확산 사전 확률의 역할을 확인시켜 준다.
- CUSP 사전 확률은 초기 계수가 심하게 과소 지정된 경우(예: rank 200 대신 10에서 시작) 복구에 결정적인 역할을 하였다.
확산 샘플러와의 비교: PnP-DM과 비교했을 때, DiffBCP는 노이즈 레벨 스케줄링에 대한 안정성이 더 높고 고해상도 이미지에서 더 나은 성능을 보였다. 다만, PnP-DM은 노이즈 파라미터를 최적값으로 수동 튜닝할 경우 균일 마스크(uniform mask)에서 경쟁력 있는 결과를 낼 수 있었다.

의의 및 주장

본 논문은 DiffBCP가 저계수 텐서 분해의 구조적 귀납 편향(inductive bias)과 현대적인 생성 확산 모델의 표현력을 성공적으로 연결했다고 주장한다. 이러한 통합을 완전한 확률적 프레임워크 내에서 정식화함으로써 저자들은 다음을 달성하였다:

노이즈가 많거나 불완전한 텐서 데이터에 대한 최첨단(SOTA) 복구 성능.
미지의 계수 및 노이즈 레벨에 대한 자동 적응, 이를 통해 광범위한 하이퍼파라미터 튜닝의 필요성을 제거하였다.
확산 사전 확률이 OOD 데이터를 마주할 때 저계수 구조를 안정화 기제로 활용함으로써, 분포 변화 및 고해상도 입력에 대한 견고성을 확보하였다.

저자들은 본 방법의 성능가 신호가 의미 있는 저계수 구조를 가진다는 가정에 의존한다고 언급한다. 만약 이 가정이 위배되면 구조적 기여도는 감소하지만, 확산 사전 확률은 여전히 활성화되어 작동한다. 향후 연구로는 이 프레임워크를 다른 역문제(예: 압축 센싱) 및 다른 형태의 텐서 분해(Tensor Train, Tensor Ring)로 확장하는 것이 제안되었다.