본문 바로가기
ML&DL/NLP(LLM)

[Prompt Engineering] Automatic Prompt Optimization with "Gradient Descent" and Beam Search Paper Review

by 거북이주인장 2024. 5. 24.

지난 포스팅에서 프롬프트 엔지니어링, 그리고 이의 확장된 개념인 프롬프트 최적화에 대해서 살펴보았다. 아직 이 글을 읽지 않았다면, 원활한 이해를 위해 먼저 읽고 오는 것을 추천한다.

https://steady-programming.tistory.com/83

 

[NLP][Prompt Engineering] What is Prompt Optimization

Prompt Engineering바야흐로 LLM의 시대이다. chatgpt를 시작으로 llama, mistral, mixtral, claude 등 다양한 llm이 쏟아져나오고 있는 가운데, 자연스럽게 프롬프트 엔지니어링 기술이 주목을 받고 있다.프롬프

steady-programming.tistory.com

한마디로 parameter space가 자연어이고 objective function에 따른 최적의 파라미터, 해를 찾는 과정이다. 본 포스팅에서는 프롬프트 최적화 기법 중 하나를 살펴볼 것이다. 해당 논문은 2023 ICLR에 accepct 되었다.

Motivation

  • manual  프롬프트 엔지니어링은 많은 수고를 필요로 하기 때문에 자동으로 해주는 절차에 대한 니즈가 자연스럽게 생겼다.
  • 기존의 프롬프트 자동 최적화 연구는 최소 low level을 통해서라도 모델에 접근할 수 있음을 가정하지만 chatgpt 등의 모델은 api을 통해서만 접근이 가능함을 가정한다.
  • 따라서 본 연구에서는 모델의 내부 weight 등은 고려하지 않고 api을 통해서만 프롬프트 최적화를 진행한다.
  • 또한 gradient descent 개념을 빌려와서 프롬프트 최적화에 적용한다.

Summary

여러 상황에서 프롬프트 최적화를 할 수 있다. openai의 프롬프트 엔지니어링 팁을 보면, 프롬프트 엔지니어링된 예시를 아래와 같이 볼 수 있다.

코드를 작성하는 프롬프트, 회의록을 요약하는 프롬프트 등 다양한 상황이 존재한다. 본 논문에서는 이 상황 중에서 분류 task에서의 프롬프트 엔지니어링에 집중한다. 즉, 더 좋은 성능을 가지기 위해 더 좋은 분류 task의 프롬프트를 어떻게 작성할 수 있을지에 대해서 살펴보는 것이다.

overview을 통해 APO의 전체적인 흐름을 파악해보자.

  • minibatch에서 initial prompt을 사용하여 prediction 값을 구한다.
  • label과 틀린 prediction을 뱉은 데이터를 추출하여 ’gradients’을 구한다.
    • llm에게 어떤 점이 부족한지 분석하게 하는 과정이다.
  • gradients을 반영하여 여러개의 새로운 프롬프트를 생성한다.
  • bandit selection을 통해서 최적의 프롬프트를 선택한다.
  • 개념 매칭
    • gradient: natural language summary of current prompt’s flaw
    • gradient descent: edit on current prompt in the opposite semantic direction of gradient
    • single gradient → multiple gradient

Methodology

Main architecture: Beam Search

개념은 생각보다 간단하다. Summary에서 살펴본 것처럼 핵심은 프롬프트의 결점 (=gradient)을 구하고 이를 반영하여 현재의 프롬프트를 수정 (=gradient descent)하는 것이다. 알고리즘은 다음과 같다.

알고리즘은 expansion, select로 구성된 beam search에 기반한다. expansion에서 양질의 후보 프롬프트를 생성하였으면 selection에서 상위 b개를 뽑는다.

  • Beam search
    • iterative 최적화 과정
    • expansion 과정과 selection 과정으로 구성됨
  • Expansion
    • 각 iteration마다 현재의 프롬프트가 새로운 후보 프롬프트를 생성하는데 사용됨
    • llm을 통해서 생성 (프롬프트)
  • Selection
    • 현재 iteration에서 후보 프롬프트를 생성했으면, 다음 iteration에 어떤 프롬프트를 가져갈지 결정하는 과정
    • 이 과정을 통해 더 나은 프롬프트 후보를 뽑는 과정
    • bandit algorithm을 통해 선택

코드를 살펴보자. 생각보다 직관적으로 짜져 있다. expansion은 `expand_candidates` 함수로, selection은 `score_candidates` 함수로 매칭된다.

Expansion step

expansion step은 다음과 같다. 

먼저 1번 과정이다. 훈련 데이터에서 minibatch을 샘플링한다. 기존의 gradient descent에서 batch gradient descent가 아닌 mini batch gradient descent을 하는 것과 유사한 원리이다. 

다음으로 2번 과정이다. minibatch의 prediction 값을 현재의 프롬프트 사용해서 구한다. 그리고 여기서 실제 label과 다른 prediction을 보이는 minibatch 데이터를 $e$에 모은다. 코드 상으로는 우선 prediction 값을 구하고 에러 데이터는 3번에서 모은다.

다음으로 3번 과정인데 여기가 조금 흥미롭다. gradient을 구한다고 하는데, 쉽게 말해서 현재 프롬프트의 결점을 분석하게 하는 것이다. 이것을 누가할까? 인간이 하기에는 너무 귀찮은 일이므로 llm에게 프롬프트를 통해서 시킨다. 코드로 살펴보자.

현재의 프롬프트가 어느 데이터에서 잘못된 결과를 내는지 {error_string}에서 알려준다. 이후에는 {num_feedbacks}개만큼의 분석을 해달라고 한다. 즉, gradient을 여러개 구하는 것이다. (multiple gradients)

논문에서 나온 예시로 살펴보자. 초기 프롬프트가 $p_0$인데 label != prediction인 에러 데이터가 $e$이다. 이제 $p_0$을 {prompt}에, $e$을 {error_string}에 넣는 것이다.

다음으로 4번을 살펴보자. 앞서 구한 multiple gradients을 사용해서 현재의 프롬프트를 업데이트한다. 눈치챘겠지만 이것도 llm에게 프롬프트를 통해서 시킨다. 코드로 살펴보자.

현재의 프롬프트가 어느 데이터에서 잘못된 결과를 내는지 뿐만 아니라 왜 그러한 결과를 내는지 분석한 내용인 gradients ( {feedback_str} )을 같이 넣는다. 그리고 {steps_per_gradient}개 만큼의 개선된 프롬프트를 만들어달라고 llm에게 시킨다.

논문의 예시로 살펴보자. 앞서 살펴본 프롬프트와 error string에 더해서 초록색 박스인 gradient가 {feedback_str}으로 들어간다.

마지막으로 수정된 프롬프트의 맥락은 유지하면서 텍스트는 서로 다른 프롬프트를 샘플링한다. 코드로 살펴보자.

107번째 라인에 프롬프트가 나와있다. 수정된 프롬프트에서 비슷한 의미의 다른 프롬프트를 샘플링하는 과정은 양질의 프롬프트를 많이 확보하고자 함에 있다.

Expansion step

다음으로 후보 프롬프트를 생성했으면 여기서 b개의 프롬프트를 선택한 후 다음 iteration으로 넘겨야 하는데, 어떻게 선택을 하는지 살펴보도록 하자. 논문에서는 프롬프트를 뽑는 과정이 Best arm identification in bandit optimization과 유사하다고 말하며 bandit 최적화 문제로 치환해서 논의를 전개한다. 여기서 잠깐 bandit optimization에 대해 살펴보고 가자.

Multi-Armed Bandits

MAB 문제로 더 유명한데, 과거 카지노에서 어느 슬롯머신에 어느 레버를 당기는 것이 최대한 많은 수익을 낼까?에서 유래되었다고 한다. 여기서 bandit은 슬롯머신, arm은 레버이다. 한정된 자원 하에서 최적의 수익을 내는 arm을 선택하는 최적화 문제이다. 한마디로 도박하는 사람들이 한정된 횟수 하에서 돈을 더 벌고 싶어서 만든 알고리즘이다.

Analogy

논문에서 이 bandit 최적화 문제와 현재 후보 프롬프트에서 b개를 선택하는 것이 아래 관점에서 유사하다고 말한다.

  • n개의 arm: n개의 프롬프트 후보
  • arm을 선택했을 때의 hidden reward: underlying dataset에 대한 프롬프트의 성능
  • pulling an arm: 랜덤하게 선택된 데이터에 대해서 프롬프트의 성능을 평가하는 것
  • reward을 최대로 하는 b개의 arm 뽑기: 성능을 최대로 뽑아내는 b개의 프롬프트 뽑기

hidden 수익을 극대화하는 arm을 찾아야하듯이, 성능을 극대화하는 프롬프트를 찾는 것이다. bandit 최적화는 여러 방법들이 있는데 논문에서는 upper confidence bound을 소개한다.

$Q_t(p_i)$는 각 프롬프트의 estimated 성능이고 $N_t(p_i)$는 각 프롬프트가 뽑힌 횟수이다. (엄격하게는 다른 의미인데 이렇게 생각하면 쉽게 이해할 수 있다.)

bandit 최적화 문제에서는 exploration과 exploitation을 진행한다. exploitation이란, 여태까지 탐색한 결과를 바탕으로 수익을 많이 내는 게임 슬롯을 집중적으로 파는 것이다. 근데 생각해보면 이렇게 소수의 게임 슬롯에 집중하면 더 큰 수익을 낼 수 있는 다른 게임 슬롯의 기회를 잡을 수 없을 것이다. 따라서 exploration을 통해 다른 게임 슬롯도 탐색한다.

UCB에서 $p_i$을 구할 때도 이 두 개념이 동일하게 사용된다. $argmax_p \{ Q_t(p) + c \sqrt{ \dfrac{\log{t}}{N_t(p)} } \}$을 살펴보자. 먼저, $Q_t(p)$은 exploitation에 해당한다. estimated 성능이 큰 프롬프트를 찾겠다는 것이니, 좋은 성능을 보이는 프롬프트에 집중하겠다는 것이다. 다음으로 $c \sqrt{ \dfrac{\log{t}}{N_t(p)} }$은 exploration에 해당한다. 해당 값은 $\log{t}$는 time step이 커질수록 , $N_t(p)$가 작아질수록 커진다. 즉, 시간이 지나도 선택을 받지 못한 ($N_t(p)$가 작은) 프롬프트에 더 큰 가중치를 둔다는 것이다. 이를 통해서 혹시나 선택받지 못한 프롬프트의 성능이 좋은지 탐색한다.

이렇게 최적화 과정을 끝냈으면 상위 b개의 프롬프트를 선택해서 다음 이터레이션으로 넘긴다.

논문에서 UCB의 단점으로 hyperparameter을 정해야함을 지적한다. $c \sqrt{ \dfrac{\log{t}}{N_t(p)} }$에서 $c$와 최적화 iteration인 $T$도 정해야한다. 이에 따라서 성능이 많이 좌지우지된다고 한다. 이에 대한 대안으로 successive rejects을 제안한다.

이 방법은 엄청 간단하다. n개의 프롬프트에 대해 n-1번의 iteration을 돌며 각 iteration에서 가장 좋지 않은 성능을 보이는 프롬프트를 하나씩 제거하는 것이다. n-1번의 iteration 후에는 가장 좋은 프롬프트만 남게될 것이다.

Experiment

score function에 따라서 분류 문제 뿐만 아니라 다른 다양한 문제에서의 프롬프트를 최적화할 수 있는데, 본 논문에서는 분류 문제에 집중한다.

Data

  • jailbreak attack detection
  • hate speech detection in online comments
  • fake news detection
  • sarcasm detection in online comments

Baselines

  • Automatic Prompt Engineering (Zhou et al., 2022)
  • GrIPS (Prasad et al., 2022)
  • TEMPERA (Zhang et al., 2023)
  • AutoGPT: open-source AI agent

Main results

  • x축: bandit algorithm 실행할 때 총 iterations 수 (한정된 자원)
  • y축: f1 score
  • iterations이 늘어날수록 f1 score도 높아지는 경향성을 보인다. (자원을 더 사용해서 좋은 프롬프트를 찾는 느낌이므로 어찌보면 당영한 결과)
  • 논문에서 제시한 방법이 baselines에 비해 좋은 성능을 보인다.

Ablation study

  • beam search와 UCB을 사용하는게 과연 좋은 결과를 내는지에 대한 ablation study이다.
  • 두 방법을 사용했을 때 좋은 성능을 보였다.

Conclusion

  • gradient 개념을 사용해서 프롬프트를 최적화하는 과정이 상당히 흥미로웠다.
  • 또한 multiple gradient을 구하고 이를 통해 얻은 수정된 프롬프트에서 샘플링하는 과정이 인상적이었다.
  • 이 모든 과정을 llm을 통해 진행함으로써 데이터만 있으면 완전한 automatic process가 가능할 것으로 보인다.
  • 다만, token 사용량 이슈, llm 대답의 variation은 함께 고려해야할 사항이지 않을까 싶다.
  • 분류 모델을 개발하고 서빙하는 비용 vs llm을 통해 분류 문제를 풀었을 때 소요되는 token 사용량을 비교하며 의사결정을 하면 되지 않을까 싶다.

댓글