Policy-Mask, 정해진 채널 예산 안에서 확산 모델 특징을 골라 분할·분류에 쓰는 기법
이반 리스토파도프(Ivan Listopadov) 등 연구진 3명이 사전공개 논문 저장소 arXiv에 올린 초록에 따르면, 이들은 가중치를 고정한 확산 모델에서 이미지 분할·분류·의미 대응 같은 판별 과제에 쓸 중간 특징을 자동으로 골라내는 기법 'Policy-Mask'를 제안했다. 논문 제목은 「Policy-Mask: Budget-Aware Adaptation of Diffusion Features to Discriminative Tasks」(arXiv:2610.28889)다. 연구자 커뮤니티 플랫폼 허깅페이스의 일일 논문 목록에서는 추천 25회를 받았다. 동료평가를 거치지 않은 사전공개 논문이다.
연구진이 풀려고 한 문제는 '어디서 특징을 꺼낼 것인가'다. 초록에 따르면 고정된 확산 모델의 중간 특징을 판별 과제에 활용하는 일은 간단하지 않다. 쓸모 있는 신호가 노이즈 제거 궤적(타임스텝), 네트워크 깊이(블록), 채널에 걸쳐 흩어져 있기 때문이다. 그래서 실제로 쓰려면 어느 타임스텝에서, 어느 블록의, 어떤 채널 묶음을 가져올지를 모두 정해야 한다. 연구진은 기존 방법들이 이 선택을 사람이 직접 고정하거나, 타임스텝처럼 일부 추출 하이퍼파라미터만 최적화하는 데 그쳤다고 지적했다. 한번 정한 규칙을 입력 이미지나 허용된 특징 크기가 달라져도 그대로 다시 쓴다는 점도 한계로 꼽았다.
Policy-Mask는 이 세 가지를 한꺼번에 고르는 정책이다. 미리 정한 채널 예산, 곧 남길 채널 수를 조건으로 받고 최종 과제의 목표 함수로 학습한다. 타임스텝·블록·채널을 각각 맡는 정책 헤드 세 개가 함께 학습되는 구조다. 연구진은 이 과정에서 생기는 불균형도 다뤘다. 별도 정규화가 없으면 선택한 채널이 많을수록 채널 점수의 기여가 커져, 타임스텝과 블록 쪽 업데이트를 압도한다. 연구진은 채널 점수를 선택한 부분집합의 크기로 나눠 정규화했다. 이렇게 세 정책 헤드의 크기를 비슷한 수준으로 맞춰 공동 학습을 안정시켰다고 설명했다.
성능은 같은 예산 조건에서 고정 추출 방식과 특징 필터 기준선을 비교해 측정했다. 초록에 따르면 Policy-Mask는 분할·분류·의미 대응 세 과제 모두에서 기준선보다 좋은 결과를 냈다. 구체적으로 장면 분할 데이터셋 ADE20K에서 38.9 mIoU, 이미지 분류 데이터셋 CIFAR-100에서 128차원 기준 top-1 정확도 75.4%를 기록했다. 이미지 간 대응점 찾기 데이터셋 SPair-71k에서는 PCK@0.1 76.3을 얻었다.
구성 요소를 하나씩 빼 보는 절제 실험(ablation)에서는 블록 선택이 최종 정확도에 실질적인 영향을 준다는 결과가 나왔다. 채널 예산에 따른 점수는 채널 256개를 남겼을 때 가장 높았다. 연구진은 이 지점을 확산 모델 특징을 압축해 쓸 때의 실용적인 작동 지점으로 제시했다. 특징을 무조건 많이 남기는 대신, 정해진 크기 안에서 어느 위치의 정보를 담을지를 학습으로 정하는 것이 이 기법의 핵심이라는 뜻이다.
연구진은 Policy-Mask를 U-Net 기반 모델에만 묶어 두지 않았다. 최근 이미지 생성에 쓰이는 확산 트랜스포머인 PixArt-Σ와 SD3.5 MMDiT에도 이 기법을 적용했다. U-Net의 채널 대신 이미지 토큰의 은닉 차원을 선택 대상으로 바꾸는 방식이다. 같은 틀로 구조가 다른 생성 모델에서도 판별 과제용 특징을 골라낼 수 있음을 보이려는 시도다.
