AMD, 이미지 생성 AI 발표했다

반도체 대기업 AMD가 자사 이미지 생성 AI인 니트로-T(Nitro-T)를 발표했다. 니트로-T는 고효율 학습에 중점을 둔 텍스트에서 이미지를 생성할 수 있는 DiT(Diffusion Transformer) 모델이다. AMD는 자사 GPU인 인스틴트(Instinct) MI300X를 32개 사용해 니트로-T를 제로부터 단 하루 미만으로 트레이닝할 수 있다고 어필하고 있다.

AMD는 2024년 11월 고속 추론에 중점을 둔 확산 모델인 니트로-1을 출시했다. AMD는 니트로-1부터 텍스트에서 이미지를 생성하는 이미지 생성 AI를 리소스 효율적인 방법으로 제로부터 트레이닝하는 방법에 대한 연구를 계속해왔다고 한다.

이 트레이닝 효율성 향상을 진행하며 AMD는 마침내 니트로-T로 32개 인스틴트 MI300X로 트레이닝을 하루 미만으로 단축하는 데 성공했다. 이로 인해 오픈소스 DiT 모델인 PixArt-α와 비교하면 트레이닝 비용을 14분의 1까지 줄일 수 있게 됐다. AMD는 니트로-T가 지닌 높은 트레이닝 효율성에 대해 스마트한 아키텍처 선택이나 시스템 최적화 등 복수 최첨단 기법을 조합해 실현했다고 설명했다.

니트로-T에는 2개 버전이 있다. 한 변이 512픽셀인 이미지 생성에 최적화된 파라미터 수가 6억 개인 DiT 모델과, 한 변이 1024픽셀인 고해상도 이미지 생성에 최적화된 파라미터 수가 12억 개인 MMDiT(Multimodal Diffusion Transformer) 모델이다. 이들 모델은 텍스트 컨디셔닝을 제공하기 위해 Llama 3.2 1B를 활용하고 있으며 트레이닝에는 패치 시퀀스 길이 단축, 수렴 고속화, 트레이닝 처리량 최적화를 가능하게 하는 전략과 설계 선택이 포함되어 있다.

트랜스포머 기반 아키텍처에서 주요 계산 병목 중 하나는 자기 어텐션 메커니즘. 자기 어텐션 메커니즘이 지닌 복잡성은 입력 토큰 수에 비례해 증대하기 때문에 트레이닝 중에 토큰 시퀀스 길이를 단축하는 건 모델 트레이닝을 고속화하는 자연스러운 방법이 된다.

확산 트랜스포머는 이미지 패치에서 생성된 토큰을 조작한다. 이미지에는 본질적으로 공간적 중복성이 존재해 토큰을 랜덤하게 드롭하거나 마스크해 토큰 시퀀스 길이를 단축할 수 있다. 하지만 단순한 토큰 드롭은 고수준 구조를 트레이닝하는 데 필요한 중요한 정보를 폐기해 모델 성능을 크게 저하시킬 수 있다.

이 문제를 해결하기 위해 MicroDiT에서 도입된 지연 마스킹 전략을 채택했다. 이 기법에서는 미리 토큰 마스크를 적용하는 대신, 트랜스포머 블록 몇 개로 구성되는 패치 믹서라고 불리는 경량 중간 모듈을 도입한다. 중요한 아이디어는 패치 믹서 단계 후까지 마스킹 동작을 지연시켜 모델이 먼저 이미지 전체 정보를 집약하고 믹스할 수 있게 하는 것. 이를 통해 토큰 서브셋이 드롭되어도 남은 토큰이 입력의 글로벌 뷰를 인코딩하고 중요한 컨텍스트와 구조 신호가 보존된다.

AMD는 지연 패치 마스킹, 딥 압축 오토인코더, 표현 얼라인먼트 등의 혁신적인 기술과 인스틴트 MI300X 상의 최신 ROCm 소프트웨어 스택을 활용해 이들 모델을 하루도 걸리지 않고 제로부터 트레이닝할 수 있음을 실증했다면서 이런 진보로 인해 연구자는 아이디어를 더 빠르게 반복할 수 있게 됐고 독립 개발자나 소규모 팀이 필요와 제약에 완벽히 맞는 모델을 트레이닝하거나 미세 조정하는 장벽이 낮아졌다면서 완전한 트레이닝 코드와 모델 가중치를 공개해 이번 연구가 실험을 촉진하고 생성 AI 도구에 대한 접근을 민주화하며 해당 분야에 대한 추가 연구 발전에 도움이 되기를 바란다고 밝혔다.

한편 니트로-T-1.2B 및 니트로-T-0.6B 모델 데이터는 허깅페이스에서 배포되고 있다. 관련 내용은 이곳에서 확인할 수 있다.