[LLM] Easy Soft prompt 개발 1탄

내가 이걸 만들게 된 계기는 관심을 가지는 분야 중 하나인 Soft prompt tuning에 대한 통합적인 코드가 없었고 베이스 모델인 T5 안에서 프롬프트가 동작하기에 메서드를 한눈에 알아보기가 어려웠다! 그래서 베이스라인 코드를 잡고 거기서 일부 코드를 수정해서 사용성을 조금 더 올라가게 만들어 보고자 개발을 했다 (사실 내 연구하려고 약간 수정한 거임).

 

기존 Soft prompt 방법들 중 대표적인 방법은 Prefix tuning (https://arxiv.org/pdf/2101.00190.pdf)으로

Prefix tuning

위와 같이 Prefix라고 하는 파인튜닝을 위한 토큰을 맨 앞에 추가함으로써 기존의 Fine-tuning보다 더 효율적으로 학습하는 것을 의미한다. 뭐.. 이런 게 시초가 되어서 인풋 임베딩 레이어에 soft prompt를 추가해서 학습하는 방법들이 나오기 시작했다. 방법들로는 Soft prompt token을 만들기 위해 LSTM 모델을 추가로 사용하는 방법, Encoder-Decoder구조에서 residual connection을 사용한 Residual prompt, 하나의 prompt token을 Low rank로 decompose 하는 방법 등등이 존재한다. 사실 생각보다 몇 개 없다!

 

 

본론으로 들어가서, 

내가 밑바닥부터 짜는 것은 매우 비효율적이며, 사실은 내 연구용으로 짠 코드라서 가장 먼저 읽은 "Decomposed Prompt Tuning via Low-Rank Reparameterization"의 코드를 기반으로 만들었다. 이 코드를 먼저 뜯어볼까?

 

Soft prompt에서 가장 먼저 봐야 하는 부분은 모델의 Embedding Layer이다. 입력으로 주어진 Text는 모델의 Embedding Layer를 통과해서 Token의 형태로 변환되는데, 이 출력에 soft prompt가 추가되기 때문이다.

 

코드를 확인해 보니 T5 모델이 있었고 크게 Encoder, Decoder로 구성이 되어있었다. Encoder에서 Embedding 하는 부분을 찾아보면,

327번 줄을 보면 input_ids가 self.embed_tokens()의 입력으로 들어가면서 임베딩된 값이 나오고, 

아래에서 self.get_prompt_encoder()를 통해 prompts_encoder를 출력하고 있다.

핵심이 되는 부분은 330줄에서 prompts_encoder와 임베딩된 값을 합치고 있다. 아! 이 prompts_encoder = soft prompt구나?

 

그럼 self.get_prompt_encoder()로 가보자.

self.prefix_first와 self.prefix_second를 행렬 곱하고 이를 batch size만큼 늘려준다. 그럼 여기서 soft prompt가 self.prefix_first와 self.prefix_second겠지? 그럼 초기화하는 부분을 가보면 학습 가능한 파라미터로 선언되어 있을 것이다.

아래와 같이!

 

예상한 것과 같이 학습 가능한 파라미터로 선언되어 있다. 아! 참고로 오리지널에서는 한 개만 선언하는데, 레퍼런스 하는 코드는 이렇게 2개를 선언한다. 그리고 또 확인해야 할 부분은 Soft prompt는 기존의 모델은 Freeze 하고 선언한 Soft prompt만 학습하기에 아래와 같이 기존 모델은 Freeze 되었는지를 확인!

 

 

근데 이렇게 되면 다른 메서드들은 적용하기가 어렵다. 왜냐하면 Freeze 하는 코드는 다른 파일에 존재하고 여기서는 'prefix'라는 변수명으로 학습할 건지를 결정하기 때문에 내 메서드에 추가되는 변수들은 또 새로 선언을 해줘야 한다. 

그래서 구조를 바꾸고 다른 메서드들도 추가해서 실험을 하고자 했다.

 

그럼 구조적으로 어떻게 바꿔볼까?

일단 내가 이루고자 하는 건 아래와 같았다.

 

1. Method들이 자신들의 코드를 자유자재로 추가할 수 있어야 함

2. 각 method들에 대한 하이퍼 파라미터를 설정하는 게 쉬워야 함

3. 모델을 Freeze 할 때, 일일이 변수명을 추가해 줄 필요가 없어야 함

 

어떻게 수정했는지는 2탄에서 알아보자~!