SegMoE: Segmind 확산 전문가 혼합

SegMoE는 사전 학습된 모델을 결합하여 처음부터 Mixture-of-Experts (MoE) Diffusion 모델을 만들기 위해 설계된 프레임워크입니다. 사용자는 Feed-Forward 블록, 어텐션 레이어 또는 전체와 같은 특정 레이어를 라우터 네트워크를 사용해 토큰을 가장 적합한 전문가에게 효율적으로 할당하는 희소 MoE 레이어로 교체할 수 있습니다.

기술 아키텍처 및 명명 규칙

SegMoE 모델은 Stable Diffusion과 동일한 기본 아키텍처를 유지하면서 여러 전문가 모델을 하나의 프레임워크에 통합합니다. 명명 규칙 SegMoE-AxB는 모델 구성을 정의합니다:

  • A: 함께 병합된 전문가 모델의 총 수.
  • B: 각 이미지를 생성할 때 활성화되는 전문가 수.

구성에 따라 특정 레이어(Feed-Forward 블록, 어텐션 또는 둘 다)만 복제되며, 나머지 파라미터는 표준 Stable Diffusion 모델과 동일하게 유지됩니다.

사용 가능한 모델 릴리스

Apache 2.0 라이선스 하에 Hugging Face Hub에 세 개의 사전 병합 모델이 공개되었습니다:

  1. SegMoE 2x1: 두 개의 전문가 모델을 사용합니다.
  2. SegMoE 4x2: 네 개의 전문가 모델을 사용합니다.
  3. SegMoE SD 4x2: 네 개의 Stable Diffusion 1.5 전문가 모델을 사용합니다.

비교 테스트 결과, SegMoE 모델은 RealVisXL_V3.0과 같은 기본 모델에 비해 프롬프트 이해도가 향상된 것으로 나타났으며, 특히 여러 객체를 렌더링하는 복잡한 상황(예: "three green glass bottles")이나 특정 공간 배치(예: "the statue of Liberty next to the Washington Monument")에서 그렇습니다.

구현 및 커스터마이징

segmoe 패키지를 사용하여 사용자 정의 MoE 모델을 만들 수 있으며, 이 패키지는 mergekit 라이브러리에서 영감을 받았습니다. 이 과정은 기본 모델 경로, 전문가 수, 혼합할 레이어 유형(ff, attn, 또는 all), 각 전문가에 대한 소스 모델 및 게이트 가중치를 계산하기 위한 해당 긍정 및 부정 프롬프트를 지정하는 config.yaml 파일을 정의하는 것을 포함합니다.

Hugging Face Diffusers와의 통합

SegMoE는 Hugging Face 생태계와 통합되어 있습니다. 추론은 segmoe 라이브러리의 SegMoEPipeline을 사용하여 수행할 수 있습니다:

from segmoe import SegMoEPipeline

pipeline = SegMoEPipeline("segmind/SegMoE-4x2-v0", device="cuda")

prompt = "cosmic canvas, orange city background, painting of a chubby cat"
negative_prompt = "nsfw, bad quality, worse quality"
img = pipeline(
    prompt=prompt,
    negative_prompt=negative_prompt,
    height=1024,
    width=1024,
    num_inference_steps=25,
    guidance_scale=7.5,
).images[0]
img.save("image.png")

성능 트레이드오프 및 하드웨어 요구사항

SegMoE가 향상된 기능을 제공하지만, 특정 계산 오버헤드를 동반합니다:

  • 추론 속도: 토큰당 전문가 수가 1보다 클 경우, 모델은 여러 전문가 모델에 걸쳐 연산을 수행해야 하므로 단일 SD 1.5 또는 SDXL 모델보다 느립니다.
  • VRAM 사용량: MoE 모델은 상당한 VRAM을 필요로 합니다. 예를 들어, SegMoE-4x2는 반정밀도에서 24GB VRAM을 요구하므로 로컬 환경보다 다중 GPU 배포에 더 적합합니다.

도구 및 배포

segmoe 도구를 통해 생성된 사용자 정의 모델은 로컬에 저장하거나 huggingface-cli 또는 huggingface_hub 라이브러리의 upload_folder 함수를 사용해 Hugging Face Hub에 푸시할 수 있습니다.

Sources