[평범한 학부생이 하는 논문 리뷰] FireFlow: Fast Inversion of Rectified Flow for Image Semantic Editing (arXiv 2412)

Paper : https://arxiv.org/abs/2412.07517

FireFlow: Fast Inversion of Rectified Flow for Image Semantic Editing

Abstract

 Rectified Flows(ReFlows)에서 image를 structured noise로 바꾸는 빠른 inversion은 아직 해결되지 않았다. 본 논문은 간단하지만 효과적인 zero-shot approach인 FireFlow를 제안한다. 이는 기존 ReFlow-based model의 능력은 그대로 가지면서, 8 step만으로 정확한 inversion과 editing하는 능력을 확장한다.

1. Introduction

 ReFlow 모델은 well-trained Reflow model이 data distribution에 걸쳐 거의 일정한 velocity dynamics를 학습하고 이는 안정성과 bounded velocity approximation error를 보장한다는 이점을 가진다. 기존의 방식은 이러한 특징을 제대로 이용하지 못하고 있다. 최근의 접근법은 정확성을 희생하고 각 step의 computational efficiency를 우선시하거나 높은 정밀도를 달성하기 위해 추가적인 cost를 발생하는 일반적인 Euler solver에 의존한다. 이에 대한 결과로, 빠르고 정확한 inversion을 하는 ReFlow Model의 potential이 아직 개발되지 않았다.

 본 논문에서는, ReFlow model을 기반으로하는 ODEs를 위한 새로운 numerical solver를 도입한다. 본 논문의 방식은 first-order solver의 computational cost를 유지하면서 second-order precision을 달성한다. 중간의 velocity approximation을 다시 사용함으로써, inversion process를 안정화하고 well-trained ReFlow model의 constant velocity property를 충분히 활용한다.

2. Motivation

 ReFlow model은 $X_t$가 $X_0$와 $X_1$사이에서 linear하게 변한다는 가정하에 작동하는데, 이는 uniform linear motion에 상응한다. 이는 다음과 같이 acceleration term을 통합하는 것으로 확장된다.

여기서 $X_{t+1} = X_t + v_t \Delta t + \frac{1}{2}a_t\Delta t^2$이다. (ChatGPT왈 : 고전 역학에서 등가속도 운동을 설명하는 기본적인 위치 업데이트 공식을 그대로 적용한 것) 최근 연구들은 (6)을 풀기 위한 training-based strategy는 coupling preservation과 inversion을 향상시킨다는 것을 보인다. 게다가 training-free method는 효과적인 inversion을 달성하는데 $v$의 2차 미분의 utility를 보였고, 이는 accelerated motion의 원칙과 맞는다.

 그러나, 이 관찰은 counterintuitive한 것으로 보인다. FLUX와 같은 well-trained ReFlow model은 $v_t$가 $X_1 - X_0$인 constant value를 근사한다고 일반적으로 가정한다. 그러므로 acceleration term $a_t = dv_t / dt$는 이론적으로 0에 근접한다.

Connection to High-Order ODE Solvers

 $a_t$를 continuous term으로 다루는 대신, 본 논문은 high-order ODE solvers의 관점으로 다시 해석한다. Finite-difference approximation $a_t = (v_{t + \Delta t} - v_t)/\Delta t$를 사용해, 다음과 같이 방정식을 다시 쓴다.

이는 second-order Runge-Kutta method의 standard formulation에 상응한다. 이 high-order approach는 p-th order method의 global error가 $\mathcal{O}(\Delta t^p)$로 scale되기 때문에 Euler의 방식과 동일한 accuracy를 달성하는데 더 적은 step(larger step size $\Delta t$)를 가능하게 한다. 이는 동일한 error tolerance $\epsilon$을 유지하면서 더 큰 $\Delta t$를 가능하게 한다. 유사하게, $a_t = (v_{t + \frac{1}{2}\Delta t} - v_t)/(\frac{1}{2} \Delta t)$를 사용해서 $a_t$를 근사하면, 결과 position update는 다음과 같다.

이는 또다른 second-order ODE solver인 standard midpoint method와 대응된다.

Impact on ReFlow Inversion

 Forward process와 동일하게 inversion과 reconstruction task에서 high-order solver는 향상된 performance를 보여준다. 본 논문에서 이러한 특징을 다음과 같이 공식화한다.

Implication

 p-th order numerical solution동안 도입된 inversion error $\Vert \Delta_T \Vert$는 reverse process에 전파되는데, 이는 학습된 drift $v_\theta(X_t, t)$의 Lipschitz 상수 $L$에 의해 scale된 약간의 reduction을 경험한다. 이러한 reduction에도 불구하고, 원래 이미지에 대한 전체적인 reconstruction error $\Vert \Delta_0 \Vert$는 동일한 order $\mathcal{O}(\Delta t^p)$로 남는다. 결국, high-order solver는 더 적은 step으로 정확한 inversion과 editing을 달성하기 위해서 ReFlow에서 선호된다.

3. Method

Challenge with High-Order Solvers

 High-order solver의 사용은 이론적으로 유망하지만, 이는 실용적인 runtime speedup을 생성하는데 실패한다. High-order solver는 interval $[t, t+1]$사이에서 더 많은 point들을 평가하는 것이 필요로 하기에 더 높은 스텝당 NFE을 야기한다. 이는 스텝 수의 감소를 부정하고 전체적인 computational efficiency를 향상에 실패한다.

 Midpoint method는 $\mathcal{O}(\Delta t^3)$의 local error와 $\mathcal{O}(\Delta t^2)$의 global error를 달성한다.

Midpoint method : 

이 방식은 스텝당 2 NFE를 필요로 한다. 한번은 $X_{t + \frac{\Delta t}{2}}$를 계산하고 다른 한번은 $v_\theta(X_{t + \frac{\Delta t}{2}}, t + \frac{\Delta t}{2})$를 계산한다. 이는 Euler method와 비교해서 cost를 두배한다. 이 midpoint method는 $v_{t + \frac{\Delta t}{2}}$를 활용해서 $v_t$보다 $\frac{X_{t+1} - X_t}{\Delta t}$에 대한 더 정확한 estimate을 제공한다. 더 적은 computational cost를 가진 대안을 찾을 필요가 있다.

A Low-Cost Alternative

 ReFlow의 training objective는 모든 $t$에 대해서 well-trained model은 $v_\theta(X_t, t) \approx (X_1 - X_0)$를 만족한다는 것을 내포한다. 이 특징을 활용해, 가장 efficient한 접근법은 $v_t$를 $v_0$으로 대체하는 것이다. 그러나 이 simplification은 multi-step iteration이 더 이상 필요하지 않기 때문에 conditional prior를 통합하는 것을 어렵게 한다.

 그래서 본 논문은 $v_t$를 $v_{t + \frac{\Delta t}{2}}$ 대신에 $v_{(t-1) + \frac{\Delta t}{2}}$로 대체하는 것을 제안한다.

Modified midpoint method : 

이 scheme에서는 스텝당 한번의 NFE만을 필요로 한다. 이는 Euler method의 computational cost와 일치한다. 이 방식의 key question은 이 scheme이 original midpoint method의 second-order accuracy를 유지하는 지이다.

 Local과 global truncation error에 대해, 본 논문은 만약 $v_\theta(X_t, t)$가 well-trained이고 $X, t$에 대해서 smooth하게 변한다면 제안된 scheme은 standard midpoint method와 동일한 truncation error를 달성한다는 것을 보인다. 이는 수정된 방식이 first-order solver의 computational cost로 계산되지만 second-order accuracy의 이점을 유지한다는 것을 보장한다.

 두 조건이 만족되면, 본 논문의 modified midpoint method는 이러한 환경 하에서 standard midpoint method와 동일한 truncation error를 달성한다. 결국, first-order Euler method와 동일한 runtime cost를 유지하면서 더 작은 전체 error를 보인다.

본 논문의 방식에 의해 생성된 transport trajectory가 더 straight하고 이는 Euler method와 동일한 NFE를 유지하면서 향상된 accuracy를 야기한다. 그리고 standard midpoint method의 성능을 능가한다.

Numerical Results and Discussion

 Approximation error는 step size $C \cdot \Delta t$와 거의 linear하게 자란다는 key trend를 따르는 것으로 보인다. 이는 이론적인 bound $\mathcal{O}(\Delta t)$와 일관된다. Error의 내재적인 변동성에도 불구하고, error의 정도는 대부분의 step에서 well-controll되고 안정적으로 보인다. 이는 재사용된 velocity approximation의 robustness를 유효하다는 것을 보인다.

Image Semantic Editing

 본 논문에서는 (https://arxiv.org/abs/2411.04746) 이 논문의 방식을 사용했다고 한다. (기회가 되면 읽어봐야겠다!) 이 논문에서는 denoising process에서 self-attention layer안의 value feature가 inversion process과정에서 생성된 value feature들로 대체되는 방식을 이용한다고 한다. 본 논문의 numerical solver의 image preservation을 활용하기 때문에, 원래 paper에서의 self-attention layer에서의 replacement를 적용할 특정한 block이나 time steps에 대한 신중한 선택을 필요로 하지 않는다. 대신에, 본 논문은 첫번째 denoising step에서만 모든 self-attention layer에 이 전략을 적용한다.

4. Experiments

4.1 Text-to-image Generation

본 논문의 method가 comparable text alignment performance를 유지하면서 더 좋은 image quality를 보인다.

4.3 Inversion and Reconstruction

Quantitative Comparison

 본 논문의 방식이 동일한 수의 step에서나 동일한 computational cost에서 상당한 reconstruction error에서의 reduction을 보인다.

Qualitative Comparison

 Efficient하고 effective한 reconstruction을 보인다.

Convergence Rate

4.4 Inversion-based Semantic Image Editing

Quantitative Comparison

 본 논문의 방식은 특히 CLIP similarity에서 다른 방식들과 비슷하거나 outperform하는 것을 볼 수 있다. 특히, 본 논문의 방식은 상대적으로 적은 step으로 고퀄리티 결과를 달성한다. 이는 이 방식의 의도된 수정과 잘 align하는 edit을 생성하면서 original content의 integrity를 유지하는데에서 효율성과 효과를 보인다.

Qualitative Comparison

Inference Speed

A. Technical Proofs

A.1 Proof of Proposition 3.1

A.2 Proof of Proposition 4.1

A.3 Proof of Theorem 4.2

B. Limitations

 본 논문의 방식이 object color를 바꾸거나 natural image에서 흔하지 않은 시나리오를 포함하는 task를 editing하는 것에 어려움을 겪는다. 흔하지 않은 discription의 경우도 잘 생성되지 않는다.

 본 논문에서는 self-attention module에서 V feature를 대체하기만 하는 editing strategy의 간단성 때문이라고 지적한다. 본 논문에서는 self attention module에서 K feature를 더하는 것을 포함하는 것이 이러한 문제를 해결할 수 있다는 것을 발견한다.

하지만, 이는 original structure와 background detail들의 보존을 감소하는 문제를 지닌다.