Precision autotuning for linear solvers via contextual bandit-based RL

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 대규모 GPU 컴퓨팅과 LLM 의 발전으로 인해 부동소수점 연산 (Floating-point arithmetic) 을 사용하는 계산 작업의 에너지 소비와 데이터 전송량이 급증하고 있습니다. 이를 해결하기 위해 저정밀도 (Low-precision, 예: FP16, BF16) 연산을 활용하는 것이 필수적이지만, 정밀도 감소는 반올림 오차 증가와 수치적 불안정성을 초래하여 계산 결과의 신뢰성을 떨어뜨릴 수 있습니다.
• 문제: 기존 혼합 정밀도 (Mixed-precision) 알고리즘은 이론적 분석이나 시행착오 (Trial-and-error) 기반의 자동 튜닝 도구에 의존합니다.
  • 이론적 분석은 조건수 (Condition number) 등 문제에 대한 사전 지식이 필요하며, 실제 적용 범위가 제한적입니다.
  • 기존 자동 튜닝 도구 (예: Precimonious, HiFPTuner) 는 컴파일/런타임 시 반복적인 실행을 통해 최적의 정밀도를 찾지만, 대규모 프로그램에서는 확장성이 떨어지고 계산 비용이 많이 듭니다. 또한, 훈련 데이터와 다른 새로운 (Unseen) 데이터에 대한 일반화 능력이 검증되지 않았습니다.
• 목표: 선형 시스템 ($Ax=b$) 을 풀기 위한 선형 솔버 (Linear Solvers) 에 대해, 강화 학습 (RL) 을 기반으로 하여 계산 단계별로 최적의 정밀도를 동적으로 선택하고, 정확도와 계산 효율성 사이의 균형을 자동으로 조정하는 프레임워크를 개발하는 것입니다.

2. 제안된 방법론 (Methodology)

저자들은 정밀도 선택 문제를 컨텍스트 밴딧 (Contextual Bandit) 문제로 공식화하고, 이를 강화 학습을 통해 해결하는 프레임워크를 제안합니다.

가. 컨텍스트 밴딧 공식화 (Contextual Bandit Formulation)

• 상태 (Context, $s$): 문제 인스턴스 (선형 시스템) 의 특성을 나타내는 특징 벡터입니다. 주로 조건수 ($\kappa(A)$) 와 행렬 노름 ($\|A\|_\infty$) 의 로그 값을 사용하여 이산화 (Discretization) 된 상태 공간으로 매핑됩니다.
• 행동 (Action, $a$): 계산 단계별 정밀도 구성 (예: LU 분해, 잔차 계산, GMRES 내부 솔버, 해 업데이트 단계별 정밀도 선택) 의 조합입니다.
• 보상 (Reward, $r$): 정확도 (Forward/Backward error), 수렴성 (반복 횟수), 계산 비용 (정밀도 수준) 을 종합적으로 고려한 다목적 함수입니다.
  • $R = w_2 \cdot f_{\text{precision}} + w_1 \cdot f_{\text{accuracy}} - f_{\text{penalty}}$
  • 여기서 $f_{\text{penalty}}$ 는 반복 횟수 증가에 따른 패널티로, 불필요한 반복을 방지합니다.

나. 학습 및 추론 프로세스

1. 이산 상태 공간 (Discretized State Space): 연속적인 상태 공간을 유한한 빈 (Bins) 으로 나누어 Q-Table 의 크기를 관리 가능하게 만듭니다.
2. 행동 공간 축소 (Action Space Reduction): 가능한 정밀도 조합의 수가 기하급수적으로 증가하는 문제를 해결하기 위해, 정밀도 순서 제약 ($u_1 \le u_2 \le \dots \le u_k$) 을 부과하여 유효한 조합만 선택하도록 제한합니다 (예: 256 개 조합을 35 개로 축소).
3. 학습 알고리즘:
   • $\epsilon$-Greedy 전략: 초기에는 다양한 정밀도 선택을 탐색 (Exploration) 하다가, 시간이 지남에 따라 최적의 행동을 선택 (Exploitation) 하도록 $\epsilon$ 값을 감쇠시킵니다.
   • Q-학습: 벨만 방정식을 기반으로 한 증분적 행동 가치 추정 (Incremental action-value estimation) 을 사용하여 Q-Table 을 업데이트합니다.
4. 적용 사례 (Case Study): GMRES 기반 반복 정제 (GMRES-IR) 알고리즘에 적용했습니다. LU 분해, 잔차 계산, GMRES 솔버, 해 업데이트 등 4 가지 주요 계산 단계에서 정밀도를 동적으로 선택합니다.

3. 주요 기여 (Key Contributions)

1. RL 기반 정밀도 자동 튜닝 프레임워크: 수치 알고리즘의 정밀도 선택을 컨텍스트 밴딧 문제로 최초로 공식화하고, 강화 학습을 적용한 일반적인 프레임워크를 제시했습니다.
2. 미시 데이터 (Unseen Data) 에 대한 일반화 검증: 훈련 세트와 완전히 다른 조건수 및 행렬 구조를 가진 테스트 데이터에 대해서도 효과적으로 작동함을 실험을 통해 입증했습니다. 이는 기존 도구들이 가지지 못한 핵심 강점입니다.
3. 지식 기반의 특징 추출: 행렬의 조건수와 노름과 같은 계산 비용이 낮은 특징을 사용하여, 복잡한 모델 학습 없이도 효율적인 정밀도 선택이 가능함을 보였습니다.
4. 다목적 보상 설계: 정확도와 계산 효율성 사이의 균형을 맞추기 위해 반복 횟수 패널티를 포함한 보상 함수를 설계하여, 수렴 실패를 방지하면서도 저정밀도 사용을 극대화했습니다.

4. 실험 결과 (Results)

실험은 밀집 행렬 (Dense Systems) 과 희소 행렬 (Sparse Systems) 에 대해 수행되었으며, BF16, TF32, FP32, FP64 정밀도를 사용했습니다.

• 밀집 행렬 (Dense Systems):

  • 조건수가 낮은 경우 (Well-conditioned): RL 에이전트는 BF16 및 TF32 와 같은 저정밀도를 적극적으로 사용하여 계산 비용을 절감하면서도 FP64 기준과 유사한 정확도를 유지했습니다. (특히 'W2' 정책에서 저정밀도 사용 빈도가 높음)
  • 조건수가 높은 경우 (Ill-conditioned): 수치적 안정성을 위해 자동으로 FP64 로 전환하여 정확도를 보장했습니다.
  • 성능: FP64 기준과 비교하여 유사한 오차 수준을 유지하면서, 조건수가 낮은 구간에서 GMRES 내부 반복 횟수를 증가시키는 대신 저정밀도 연산으로 전체 효율성을 높였습니다.

• 희소 행렬 (Sparse Systems):

  • 매우 나쁜 조건수 (High condition number) 를 가진 데이터셋에서는 RL 에이전트가 매우 보수적으로 행동하여 거의 모든 단계에서 FP64 를 선택했습니다.
  • 이는 저정밀도 사용이 수치적 불안정성 (Stagnation 또는 발산) 을 초래할 수 있음을 학습했기 때문이며, RL 이 문제의 난이도에 따라 정밀도 전략을 적응적으로 조절할 수 있음을 보여줍니다.

• 패널티 항의 중요성: 반복 횟수 패널티 ($f_{\text{penalty}}$) 를 제거한 실험에서는 RL 이 저정밀도를 과도하게 선택하여 정확도를 맞추기 위해 반복 횟수가 급증하는 현상이 관찰되었습니다. 이는 보상 함수 설계의 중요성을 강조합니다.

5. 의의 및 결론 (Significance)

• 과학적 컴퓨팅의 새로운 패러다임: 이 연구는 수치 알고리즘의 정밀도 튜닝을 위해 AI(강화 학습) 를 적용한 최초의 작업 중 하나로, 프로그래머의 수동 개입을 줄이고 자동화된 정밀도 조정을 가능하게 합니다.
• 실용성: 훈련 데이터가 제한된 상황에서도 새로운 데이터에 대한 추론 (Inference) 이 가능하며, 온라인 학습 루틴으로 구현되어 모델 재학습 없이도 적용 가능합니다.
• 확장성: 선형 솔버뿐만 아니라 다른 수치 알고리즘 (최소제곱법, GMRES 등) 으로도 확장 가능하여, 차세대 고성능 컴퓨팅 (HPC) 및 AI 하드웨어 환경에서 에너지 효율적인 계산 방법을 제시합니다.

요약하자면, 이 논문은 강화 학습을 통해 선형 솔버의 각 계산 단계에 필요한 정밀도를 동적으로 선택하는 지능형 시스템을 제안하며, 이는 정확도를 유지하면서 계산 비용을 획기적으로 줄일 수 있는 강력한 방법론임을 실험적으로 입증했습니다.