기술 요약: 미분 가능한 최적화를 위한 완전 1차 계층 (FFOLayer)

1. 문제 정의

미분 가능한 최적화(Differentiable optimization)는 수학적 프로그램을 미분 가능한 머신러닝 파이프라인 내에 임베딩하여, 모델의 출력이 최적화 문제를 해결함으로써 결정되는 엔드 투 엔드(end-to-end) 학습을 가능하게 하는 것을 목표로 한다. 이 분야의 주요 과제는 **확장성(scalability)**이다.

CvxpyLayer나 QPTH에서 사용되는 표준 방식들은 카루시-쿤-터커(KKT) 조건의 **음함수 미분(implicit differentiation)**에 의존한다. 이 과정은 목적 함수의 헤시안(Hessian)을 포함하는 KKT 행렬의 역행렬을 포함하는 선형 시스템을 계산하고 풀어야 한다. 대규모 문제의 경우, 이러한 행렬을 인수분해하고 역행렬을 구하는 작업은 계산 비용이 많이 들고 메모리 집약적이며, 종종 학습의 병목 현상이 된다. 또한, 헤시안을 피하려고 시도하는 기존 방법들(예: LPGD)은 유한 시간 수렴 보장이 부족하거나 점근적 가정이 필요하며, 다른 방법들(예: BPQP, dQP)은 문제를 이차 계획법(QP)으로 변환하여 오버헤드를 발생시킨다.

본 논문은 일반적인 볼록 제약 조건(convex constraints)과 호환되면서도 이론적인 수렴 보장을 유지하며, 오직 1차 정보(그레이디언트 및 함수 값)만을 사용하여 하이퍼그레이디언트(hypergradient)를 계산하는 방법에 대한 필요성을 다룬다.

2. 방법론

저자들은 FFOLayer를 제안한다. 이는 미분 가능한 최적화 문제를 바이레벨 최적화(bilevel optimization) 인스턴스로 재구성하고, 새로운 **활성 집합 라그랑주 오라클(active-set Lagrangian oracle)**을 도입한 것이다.

2.1 바이레벨 재구성

원래의 문제는 다음과 같이 상위 레벨 목적 함수 $F(x)$가 하위 레벨의 제약 최적화 문제의 해 $y^*(x)$에 의존하는 바이레벨 최적화로 간주된다:
$$\min_x f(x, y^*) \quad \text{s.t.} \quad y^* \in \arg\min_y g(x, y) \text{ s.t. } h(x,y) \le 0, e(x,y)=0$$

2.2 활성 집합 축소 ("유령" 문제)

2차 정보(헤시안)를 피하기 위해, 저자들은 "유령(ghost)" 바이레벨 문제를 구성한다.

1. 활성 집합 식별: 주어진 지점 $\bar{x}$에서, 활성 부등식 제약 조건 집합 $I$를 식별한다 (여기서 $h_i(\bar{x}, y^*) = 0$ 이고 $\lambda^*_i > 0$인 경우).
2. 선형화: 활성 부등식 제약 조건을 현재 해 주변에서 선형화하여 이를 등식 제약 조건으로 변환한다.
3. 대리 목적 함수: 하위 레벨 목적 함수를 활성 제약 조건에 대한 라그랑주 항을 포함하도록 수정하며, 이때 듀얼 변수 $\lambda^*$와 $\nu^*$를 고정된 상수로 취급한다.

이를 통해 하위 레벨 문제는 오직 선형 등식 제약 조건만을 갖는 문제로 변환된다. 저자들은 선형 독립 제약 조건 규격(LICQ)과 활성 집합의 국소적 불변성 하에서, 이 "유령" 문제의 하이퍼그레이디언트가 원래 문제와 동일함을 증명한다.

2.3 유한 차분 하이퍼그레이디언트 근사

이 방법은 KKT 시스템을 푸는 대신, 유한 차분 전략을 통해 하이퍼그레이디언트를 계산한다:

1. 섭동(Perturbation): 상위 레벨 목적 함수에 작은 섭동 $\delta$를 하위 레벨 문제에 주입한다.
2. 두 번의 풀이: 알고리즘은 블랙박스 솔버를 사용하여 원래의 하위 레벨 문제와 섭동된 "유령" 문제를 각각 푼다.
3. 근사: 두 해의 차이를 사용하여 민감도 항 $dy^*/dx$를 근사한다.

이 접근 방식은 부등식 제약 조건에 대한 페널티 기반 방법의 $O(1/\epsilon)$ 비용 대신, $\epsilon$-근사 하이퍼그레이디언트를 달성하기 위해 $O(\log(1/\epsilon))$ 번의 1차 오라클 호출만을 요구한다.

2.4 솔버 불가지론적 구현 (Solver-Agnostic Implementation)

구현의 핵심 특징은 하위 레벨 솔버를 블랙박스로 취급한다는 점이다.

• 이 방식은 솔버가 특정 미분 규칙을 지원할 필요가 없다 (예: diffcp에 의존하는 CvxpyLayer와 달리).
• GUROBI 및 MOSEK과 같은 최첨단 솔버를 수정 없이 활용할 수 있다.
• 역전파(backward pass)는 상위 레벨 그레이디언트로부터 유도된 선형화된 목적 함수를 사용하여 섭동된 문제를 푸는 방식으로 구현되며, 명시적인 헤시안 계산을 요구하지 않는다.

3. 주요 기여

1. 1차 바이레벨 오라클: 본 논문은 오직 1차 정보만을 사용하여 일반적인 제약 바이레벨 최적화에 대해 $\epsilon$-근사 하이퍼그레이디언트를 계산하는 새로운 오라클을 도입하며, 계산 복잡도는 $O(\log(1/\epsilon))$이다.
2. 수렴성 증명: 저자들은 이 방법이 $\tilde{O}(\delta^{-1}\epsilon^{-3})$의 복잡도로 $(\delta, \epsilon)$-Goldstein 정지점(stationary point)에 도달함을 보여주는 수렴 분석을 제공한다. 이는 비매끄러운 비볼록 최적화(nonsmooth nonconvex optimization)의 최상위 수준 속도와 일치하며, 선형 제약 설정에서 일반 볼록 제약 설정으로 보증 범위를 확장한다.
3. 오픈 소스 라이브리: 저자들은 기존의 미분 가능한 최적화 인터페이스(예: CvxpyLayer)와 호환되는 PyTorch 라이브러리인 FFOLayer를 공개하며, 이는 일반적인 볼록 제약과 임의의 솔버를 지원하는 드롭인 교체(drop-in replacement) 옵션을 제공한다.

4. 실험 결과

이 방법은 세 가지 태스크(합성 의사결정 중심 학습(DFL) 태스크, 스도쿠 제약 학습 태스크, 그리고 2차 원뿔 계획법(SOCP) 태스크)에 대해 평가되었다.

• 수렴성: FFOLayer(일반 볼록을 위한 FFOCP 및 QP를 위한 FFOQP 모두)는 CvxpyLayer 및 QPTH와 같은 정확한 미분 방법과 유사한 훈련 손실 수렴도를 달성했다.
• 실행 시간 효율성:
  • 역전파(Backward Pass): FFOLayer는 특히 KKT 행렬 인수분해가 불안정하고 느려지는 악조건의 문제(예: 스도쿠)에서 CvxpyLayer보다 상당한 속도 향상을 보였다.
  • 확장성: 문제 차원($y$)이 증가함에 따라, FFOLayer는 아선형(sublinear) 계산 시간 성장과 거의 일정한 피크 메모리 사용량을 나타냈다. 반면, CvxpyLayer는 초선형 시간 성장과 과도한 메모리 사용량을 보였으며, FFOLayer가 성공한 차원에서는 실행에 실패했다.
• 다른 1차 방법과의 비교: FFOLayer는 속도와 안정성 모두에서 LPGD보다 우수한 성능을 보였는데, 이는 LPGD의 점근적 접근 방식과 대비되는 FFOLayer의 명시적인 비점근적(non-asymptotic) 하이퍼그레이디언트 보장 덕분이다. 또한, 일반 문제를 QP로 변환하는 오버헤드를 피함으로써 BPQP 및 dQP보다 우수한 성능을 보였다.

5. 의의 및 주장

본 논문은 FFOLayer가 미분 가능한 최적화의 주요 병목 현상인 하이퍼그레이디언트 계산의 확장성 문제를 해결한다고 주장한다. 헤시안 계산과 KKT 행렬 역행렬 계산을 제거함으로써, 이 방법은 이전 방법들이 메모리 또는 시간 제약으로 인해 불가능했던 대규모 설정에서도 미분 가능한 최적화 계층을 사용할 수 있게 한다.

저자들은 자신들의 접근 방식이 **솔버 불가지론적(solver-agnostic)**임을 강조하며, 이를 통해 실무자들이 커스텀 미분 로직을 구현하지 않고도 최적화된 상용 솔버(예: GUROBI)를 활용할 수 있다고 밝힌다. 이론적으로, 이 연구는 1차 바이레벨 최적화와 실제적인 미분 가능 계층 사이의 간극을 메우며, 일반적인 볼록 제약, 솔버 불가지론, 그리고 유한 시간 수렴 보장을 결속시킨 첫 번째 방법이다.

결론적으로, 이 방법은 근사치를 도입하지만, 실험 결과 이러한 근사가 최적화 성능을 저하시키지 않음을 보여줌으로써 정확한 음함수 미분(exact implicit differentiation)을 대체할 수 있는 더 확장 가능한 대안임을 입증한다.