⚡ electrical engineering

AGG: Jacobian-Aggregated Group Gradient for Efficient GRPO Training of Diffusion Models

이 논문은 확산 모델 궤적의 근선형적 변화를 활용하여 중간 자코비안(Jacobian)을 근사하고 그룹 그래디언트를 집계함으로써, 디퓨전 트랜스포머(Diffusion Transformer)의 GRPO(Group Relative Policy Optimization) 학습 비용을 유의미한 품질 저하 없이 약 2배 절감하는 방법론인 JAGG(Jacobian-Aggregated Group Gradient)를 소개한다.

원저자: Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

게시일 2026-07-21

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

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

당신이 디지털 아티스트에게 더 나은 그림을 그리는 법을 가르치려 한다고 상상해 보세요. 이 아티스트는 인간이 아니라, '확산 모델(diffusion model)'이라고 불리는 복잡한 컴퓨터 프로그램입니다. 이 모델을 노이즈와 정전기 가득한 찰흙 덩어리에서 시작해, 한 단계씩 노이즈를 깎아내어 완벽한 조각상을 만들어내는 조각가라고 생각해 봅시다. 이 조각가가 실제로 인간이 좋아하는 예술품을 만들도록 가르치기 위해, 우리는 '강화 학습(Reinforcement Learning)'이라는 기술을 사용합니다. 이것은 조각가가 몇 가지 다른 조각 동작을 시도하고, 심판(보상 모델)으로부터 점수를 받은 뒤, 그 점수를 바탕으로 다음번에 더 잘하도록 배우는 게임과 같습니다.

문제는 이 학습 과정이 엄청나게 비용이 많이 든다는 점입니다. 조각가가 아주 작은 조각(하나의 '단계')을 할 때마다, 컴퓨터는 조각가의 뇌를 정확히 어떻게 조정해야 할지 알아내기 위해 거대하고 에너지를 많이 소모하는 계산을 수행해야 합니다. 만약 조각상이 완성되기까지 20단계가 필요하다면, 컴퓨터는 매번 학습할 때마다 이 무거운 수학 계산을 20번 수행해야 합니다. 이는 마치 새로운 춤을 배울 때, 동작 하나를 할 때마다 멈춰 서서 코치에게 전체적인 비평을 요구하고, 근육 기억 전체를 새로 쓰고, 다시 처음부터 시작하는 것과 같습니다. 고해상도 이미지의 경우, 이 작업은 너무 많은 시간과 전기를 소모하여 이러한 모델을 효과적으로 훈련하는 것을 거의 불가능하게 만듭니다.

여기서 JAGG라는 새로운 아이디어가 등장합니다. 이 논문의 연구자들은 영리한 질문을 던졌습니다: "우리가 정말 매 동작마다 멈춰서 모든 것을 다시 계산해야 할까?" 그들은 조각 과정의 초기, 무질서한 단계에서는 변화가 매우 부드럽고 예측 가능한 방식으로 일어난다는 점, 즉 거의 직선에 가깝다는 점을 발견했습니다. 만약 조각이 어디서 시작해서 몇 번의 움직임 후에 어디에 도달하는지 안다면, 무거운 수학 계산을 매 단계마다 하지 않고도 중간에 어떤 일이 일어났는지 실제로 추측할 수 있습니다.

이 논문은 JAGG(Jacobian-Aggregated Group Gradient)라는 방법을 소개합니다. JAGG는 모델이 한 그룹의 단계(예를 들어 네 번의 연속된 동작)를 수행하게 한 뒤, '뇌 업데이트'를 위한 무거운 수학 계산을 맨 처음 한 번, 그리고 맨 마지막에 한 번, 총 두 번만 수행합니다. 중간 단계들을 위해서는 스마트한 지름길을 사용합니다. 시작점과 끝점 사이의 경로가 직선이라고 가정하고 그 간극을 채우는 방식입니다.

연구진은 조각 과정이 완벽하게 부드럽고 선형적이라면, 이 지름길이 실제로 완벽하며, 매번 힘든 작업을 수행했을 때와 정확히 동일한 답을 준다는 것을 수학적으로 증명했습니다. 현실 세계에서는 모든 것이 완벽하게 직선은 아니지만, 그들은 이 지름길이 이미지 생성의 초기 '노이즈' 단계에서 놀라울 정도로 잘 작동한다는 것을 발견했습니다. 이 방법을 사용함으로써, 그들은 이 모델들을 훈련시키는 데 걸리는 시간을 단계당 약 17%에서 18%까지 단축할 수 있었습니다.

가장 좋은 점은 무엇일까요? 이미지의 품질이 저하되지 않았다는 것입니다. 연구진이 Flux나 QwenImage와 같은 다양한 모델에 JAGG를 테스트했을 때, 생성된 그림들은 느리고 전통적인 방식이 만든 것만큼이나 훌륭했습니다. 사실, 어떤 경우에는 이 지름길이 모델이 더 안정적으로 학습하는 데 도움을 주기도 했습니다. 이 논문은 우리가 언제 무거운 작업을 수행할지에 대해 조금 더 영리하게 생각함으로써, 최종 결과물의 아름다움을 희생하지 않고도 이 강력한 AI 아티스트들을 훨씬 더 빠르고 저렴하게 훈련할 수 있음을 보여줍니다.

기술 요약: JAGG (Jacobian-Aggregated Group Gradient)

문제 정의

Group Relative Policy Optimization (GRPO)는 생성 모델을 인간의 선호도에 맞게 정렬하기 위한 선도적인 강화 학습(RL) 알고리즘으로 부상하였으며, 대규모 언어 모델(LLM)에서 별도의 가치 네트워크(value network) 없이도 성공적으로 작동하고 있습니다. 그러나 GRPO를 확산 모델(diffusion models) 및 플로우 매칭 모델(flow matching models, 특히 Diffusion Transformers 또는 DiTs)로 확장하는 데에는 심각한 계산적 병목 현상이 발생합니다.

확산 기반 GRPO에서 생성 과정은 $W$ 개의 타임스텝(timestep)에 걸친 샘플링 궤적(trajectory)에 대해 정책을 최적화하는 다단계 결정 문제로 취급됩니다. 표준 학습 방식은 정책 그래디언트를 계산하기 위해 궤적 내의 모든 타임스텝에 대해 고용량 DiT 백본의 풀 포워드(forward) 및 백워드(backward) 패스를 수행해야 합니다. 고해상도 텍스트-이미지(T2I) 생성을 위해 궤적이 수십 단계에 이르는 경우, 이는 매우 과도한 학습 비용을 초래합니다. 토큰 단계가 가벼운 언어 모델과 달리, 현대적인 DiT를 통한 각 백워드 패스는 트랜스포머 포워드 패스의 전체 비용을 발생시키며, 이로 인해 타임스텝당 백워드 요구 사항이 지배적인 학습 비용이 됩니다.

추론 가속화를 위한 학습 불필요 방법들(예: $\Delta$ -DiT, ScalingCache)은 DiT의 은닉 상태(hidden states)와 속도 예측값(velocity predictions)이 디노이징 궤적을 따라 매끄럽고 거의 선형적으로 변화한다는 관찰 결과를 활용하여 중간 평가를 건너뛰는 데 성공했지만, 이러한 선형성을 RL 학습 시 백워드 패스 비용을 줄이는 데 활용할 수 있는지 여부는 미해결 과제로 남아 있었습니다.

방법론: JAGG

저자들은 $W$ 개의 연속적인 디노이징 스텝 그룹에 대해 네트워크 백워드 패스 횟수를 $W$ 에서 2로 줄이는 수학적으로 원칙 있는 가속화 방법인 **JAGG (Jacobian-Aggregated Group Gradient)**를 제안합니다.

핵심 통찰

이 방법은 모델 파라미터 $\theta$ 에 대한 디노이징 속도장(velocity field) $v_\theta(z, t)$ 의 자코비안(Jacobian)이 입력 $(z, t)$ , 특히 디노이징 과정의 고노이즈(high-noise) 영역에서 근사적인 선형성을 보인다는 통찰에 기반합니다. 이를 통해 양 끝점의 자코비안( $J_0$ 및 $J_{W-1}$ )만을 사용하여 중간 단계의 자코비안을 선형 보간법으로 정확하게 근사할 수 있습니다.

알고리즘 워크플로우

JAGG는 네 가지 단계를 통해 $W$ 개의 타임스텝을 처리합니다:

포워드 패스 (Forward Passes): 모델은 $W$ 개 단계에 대해 포워드 패스를 수행합니다. 그러나 그래디언트 추적은 오직 두 개의 끝점(단계 0 및 단계 $W-1$ )에 대해서만 활성화됩니다. 중간의 $W-2$ 개 단계는 메모리와 계산량을 절약하기 위해 torch.no_grad()를 사용하여 실행됩니다.
업스트림 그래디언트 계산 (Upstream Gradient Computation): 각 단계 $j$ 에 대해, 모델 출력의 분리된 프록시(detached proxy)가 생성됩니다. PPO 스타일의 목적 함수로부터 유도된 단계별 손실 $L_j$ 가 계산되며, 업스트림 그래디언트 $s_j = \partial L_j / \partial \hat{v}_j$ 가 자동 미분을 통해 계산됩니다. 결정적으로, 이 단계에는 트랜스포머 백워드 패스가 포함되지 않습니다. 그래디언트는 가벼운 SDE 단계와 PPO 목적 함수를 통해서만 흐릅니다.
타임스텝 가중 집계 (Timestep-Weighted Aggregation): 타임스텝 $t_j$ 에서 유도된 보간 가중치 $\alpha_j$ 를 사용하여 업스트림 그래디언트 $s_j$ 를 두 개의 합성 신호 $s_{base}$ 와 $s_{corr}$ 로 집계합니다:
$\alpha_j = \frac{t_0 - t_j}{t_0 - t_{W-1}}$
$s_{base} = \sum_{j=0}^{W-1} (1 - \alpha_j) s_j, \quad s_{corr} = \sum_{j=0}^{W-1} \alpha_j s_j$
이 집계는 선형 보간 가정을 통해 전체 그래디언트 신호를 효과적으로 재구성합니다.
조인트 백워드 (Joint Backward): 두 개의 끝점 출력 $[\hat{v}_0, \hat{v}_{W-1}]$ 에 대해 집계된 그래디언트 $[s_{base}, s_{corr}]$ 를 사용하여 단 한 번의 torch.autograd.backward 호출이 실행됩니다. 이는 표준 학습에서 요구되는 $W$ 번의 개별 백워드 패스를 대체하여, 모듈당 정확히 한 번의 FSDP all-gather를 트리거합니다.

이론적 보장

저자들은 $v_\theta(z, t)$ 가 $(z, t)$ 에 대해 엄격하게 선형인 경우, $t$ -가중 자코비안 보간이 근사 오차 없이 **정확(exact)**하다는 것을 명시하는 공식 증명(Theorem 1)을 제공합니다. 이러한 조건 하에서, 두 끝점 백워드 패스는 정확한 그룹 그래디언트를 복구합니다.

적응형 라우팅 ( $jagg\_frac$ )

저자들은 디노이징의 저노이즈(후기) 단계에서 선형성 가정이 저하된다는 점을 인식하여, 라우팅 파라미터 $jagg\_frac$ 을 도입했습니다. 이 파라미터는 고노이즈 선형성 가정이 성립하는 것으로 경험적으로 검증된(그래디언트의 코사인 유사도를 통해 확인) 첫 번째 일부 디노이징 그룹에만 JAGG를 적용하고, 나머지 그룹에는 정확한 단계별 백워드 패스를 사용합니다. 테스트된 T2I 모델들의 경우, $jagg\_frac$ 은 0.6으로 설정되었습니다.

주요 기여

병목 현상 식별: 본 논문은 샘플링 궤적 전반에 걸친 단계별 백워드 패스로 인해 발생하는 확산 GRPO 학습의 계산적 병목 현상을 공식적으로 식별하고 규명했습니다.
JAGG 프레임워크: 속도 자코비안 선형성에 대한 엄밀한 이론적 분석에 근거하여, 백워드 패스 횟수를 그룹당 $W$ 에서 2로 줄이는 자코비안 집계 프레임워크를 제안했습니다.
이론적 증명: 저자들은 제안된 보간법이 선형 속도 조건 하에서 정확함을 입증하여, 근사의 견고한 이론적 토대를 제공했습니다.
경험적 검증: 다양한 확산 백본(Flux, QwenImage)과 RL 알고리즘(DanceGRPO, Flow-GRPO-Fast)에 대한 광범-한 실험을 통해 방법의 일반성과 효과를 입증했습니다.

실험 결과

저자들은 HPSv2.1, PickScore, CLIPScore를 지표로 사용하여 T2I 벤치마크에서 JAGG를 평가했습니다.

효율성: JAGG는 표준 DanceGRPO 설정에서 학습 단계당 **17.3% ~ 18.2%**의 일관된 실제 시간(wall-clock) 가속을 달성했습니다(QwenImage의 경우 약 835초에서 691초로, Flux의 경우 약 279초에서 228초로 단축). Flow-GRPO-Fast에 적용했을 때, 업데이트 단계 자체에서 **29.9% ~ 30.0%**의 가속을 보였습니다.
품질: 이 방법은 생성 품질의 무시할 만한 저하와 함께 이러한 가속을 실현했습니다. 많은 경우, JAGG는 베이스라인 보상 곡선을 밀접하게 추적하였으며 수렴 시 베이스라인과 일치하거나 약간 상회하는 성능을 보였습니다.
절제 연구 (Ablation): 끝점에서는 백워드 패스를 수행하지만 자코비안 보간 없이(끝점에 독립적으로 전체 손실을 할당) 수행한 변형 모델은 일관되게 성능이 낮았습니다. 이는 자코비안 보간을 통한 중간 그래디언트 정보의 집계가 안정적인 학습에 결정적임을 확인시켜 줍니다.
그래디언트 품질: 경험적 분석 결과, 고노이즈 그룹(0–2)에서 JAGG로 집계된 그래디언트와 정확한 그래디언트 사이의 코사인 유사도가 높게(0.74–0.85) 나타나 선형성 가정을 검증했습니다. 이후 그룹에서는 유사도가 떨어졌으며, 이는 $jagg\_frac$ 라우팅 전략의 정당성을 뒷받침합니다.

의의 및 주장

본 논문은 JAGG가 대규모 확산 GRPO 학습을 실용화하는 데 중요한 진전이라고 주장합니다. 고노이즈 영역에서의 속도 자코비안의 근사적 선형성을 활용함으로써, JAGG는 생성 품질을 희생하지 않고 RL 학습의 계산 부담을 줄입니다.

저자들은 JAGG가 단순한 휴리스틱 숏컷이 아니라 디노이징 궤적의 수학적 특성에 근거하고 있음을 강조합니다. 이 방법은 기존의 추론 가속 기술과는 구별되는, 확산 RL 특유의 백워드 패스 병목 현상을 타겟팅하는 것으로 제시됩니다. 결과적으로 JAGG는 학습 비용 장벽을 크게 낮춤으로써 시각적 생성 모델을 위한 RL 기반 정렬의 광범위한 채택을 가능하게 할 것입니다.