MedAI) LLM-based Medical Assistant Personalization with Short- and Long-Term Memory Coordination

사실 research problem을 정확히 이해한 건지 잘 모르겠긴 한데..

 

personalization에 관련된 논문입니다. 아무래도 의료계니까 더욱 중요성이 부각되겠지요~.~ 

LLM의 output을 개선하기 위해, retriever를 활용해 지난 세션에서 LLM이 냈던 에러와 그에 대한 유저 피드백을 메모리에서 불러와 현재 세션의 응답을 개선하는, 뭐 그런 시도들을 했었나봐요. 저자는 이런 방식의 두 가지 문제점을 지적합니다.

1. 대부분 이런 메모리는 key-value 구조인 dictionary 형식이어서 retriever의 성능에 크게 의존한대요. key: previous mistake, value: corresponding user-feedback이라고 말하는데 이 구조 자체가 왜 문제라고 지적하는진 잘 모르겠네요 hmm
2. 학습 없이는 personalization이 쉽지 않은 방식이라고 합니다. 예시를 들어주는데, 어떤 환자는 알잘딱깔센 처방을 바라고, 어떤 환자는 여기에 구구절절 자세한 설명을 보태주길 바라는데, 즉 처방 스타일도 취향이 갈리는데 이게 반영이 안된다고 하는 거예요. 여기서 하고 싶은 말은, LLM augment해서 정확성 높이는 건 좋은데, personalization까진 한계가 있다~ 정도인듯

저자들은 PEFT + 메모리 메커니즘을 제시합니다.

메모리 메커니즘을 먼저 간략히 소개하는데요, 우리의 기억이 Dual process로 진행이 된다고 해요: 무의식적인 것과 의식적인 것. 우리가 무언가에 숙달되면 (이를테면 타자 400타 이상 나오는..) 그게 곧 무의식적으로 변한다고 하는데, 여기서 착안점을 얻었나봅니다. 이들이 제안하는 DPeM (Dual Process enhanced Memory)는 working memory (작업기억), short-term memory(단기기억), long-term memory (장기기억) 이렇게 세 종류의 메모리를 종합적으로 활용하는 건가봐요. 

PEFT는 뭐.. 새로운 건 아니니까 ㅇㅅㅇa

 

Methodology

PEFT를 활용해 의료 도메인 특화의 adapter를 붙여 훈련을 하는데요, 2020년에 Gururangan et al.이 발표한 연구에 따르면 이렇게 adapter로 일부 파라미터만 업데이트할 경우 소위 '망각' 문제가 생긴대요. 이 현상에 대한 대책으로 sample loss를 정의해서 사용을 하는데, 아래와 같은 식입니다.

$$\begin{gather} \mathcal{L}_K = -\dfrac{1}{K}\sum_i\log p(m_i) \label{eq1}\tag{1} \\ \mathcal{L}_S = \Vert V_o, V_k \Vert^2_2 \label{eq2}\tag{2}\end{gather}$$

$\ref{eq1}$은 knowledge loss인데, 의료분야 텍스트에서 $K$개의 masked token이 있다고 했을 때 $p(m_i)$는 $m_i$를 생성할 확률입니다. 한편, $\ref{eq2}$는 $V_o$와 $V_k$ 사이의 L2 distance인데요, adapter 끼우기 전 원래 layer에서의 마스크 처리된 토큰의 벡터와 어댑터 끼운 layer의 벡터 간 거리입니다. 합리적인 생각인 것 같음!

 

이들이 제시하는 DPeM은 Rehearsal과 Executive 두 개로 구성된 dual process입니다. 

1. 리허설(Rehearsal)에서는 메디컬 의사-환자 대화 스크립트들 모음집인 $\mathcal{D}$로 학습을 하고 (learning), 여기서 학습된 정보들이 작업기억에 저장이 됩니다. 우리가 새로운 정보를 알게 되면 note-taking을 하잖아요? 그런 note를 생성하는 식으로 메모리에 저장하나봄ㅋㅋ 그리고 이 작업기억은 $\mathcal{D}$의 문서들을 학습함에 따라 계속 새로고침(?)된다고 합니다. 이후에 요약(summarize) 과정에서는 이런 노트들 중에 'relevant'하다고 판단된 애들이 STM으로 이전됩니다. 이 과정에는 coordinator라 불리는 agent —이 용어를 써도 된다면— 을 써요.
2. 이렇게 작업기억에 저장된 지식들은 장기기억에 저장될지 단기기억에 저장될지 실행(Executive) 단계에서 결정된다고 해요. 이걸 보고 기억(memorizing) 과정이라고 부르는 거 같네요! 일단 작업기억으로부터 필터링되어 넘어온 지식 $k$는 상식(common-sense) 또는 사용자 특화(user-specific) 정보 이렇게 두 종류로 분류되고, STM에는 종류(key)-지식(value)의 dictionary 형태로 저장됩니다. LTM으로 넘어가는 빈도의 판단 기준은 그 정보를 불러오는 빈도인데, hyperparameter인 $\theta$보다 빈번하면 LTM으로 넘기는 거 가틈.. 정보 회수의 빈도는 학습 과정에서 flag variable로 카운트하는 거 같아요.

작업기억은 단기기억을 위한 버퍼(buffer) 정도의 역할을 한다고 이해하면 될 거 같아요.  어쨌든 STM과 LTM은 look-up이 가능한 형태라고 합니다. 주목할 만한 점이라고 하면.. 작업기억과 STM은 제한된 저장공간을 갖고 있는 반면에 LTM은 무한하다는 것? 뇌를 너무 과대평가하는 거 아닌지

 

최종적으로 기억저장소 $\mathcal{M}$은 아래와 같이 구성됩니다.

$$\begin{align}\label{eq3}\tag{3} \mathcal{M}_{\text{working}} & = \{ nt_0,\: \cdots, \:nt_i,\: \cdots\}, \\ \mathcal{M}_{\text{STM}} & = \{ \cdots, \: k\_\text{type}:k_j,\: \cdots\}, \\ \mathcal{M}_\text{LTM} & = \{ \cdots, \: k\_\text{type}:k_f,\: \cdots\}, \\ \mathcal{M} & = \{ \mathcal{M}_{\text{working}},\: \mathcal{M}_{\text{STM}}, \:\mathcal{M}_\text{LTM}  \}  \end{align}$$

 

사용자가 대화를 나눌 LLM을 학습시킬 때는 동일한 dialogue set을 활용을 하는데, 이때 PEFT를 쓴대요. 최종적으로 inference 단계에서 retriever를 사용해 memory에서 필요한 dialogue 정보를 불러오는 모양인데요, LTM과 STM에 각각 다른 retriever를 적용합니다. $\mathcal{R}_\mathcal{C}$는 STM 전용 retriever인데, 얘는 Levenshtein distance —문자열 s를 s'으로 변형하기 위해 필요한 삭제, 삽입, 교체의 최소 횟수— 기준 가장 query와 유사한 지식을 불러오고, $\mathcal{R}_\mathcal{S}$는 우리가 잘 아는 코사인 유사도 기반의 retriever입니다.

 

아따 복잡하네.. proceeding에 실릴 만했다..

 

Dataset

dialogue 데이터셋 만들 때는 챗봇이랑 환자가 실제로 대화한 거 + self-chat을 썼다고 하구요, 여기에 의사가 검토작업에 들어간 모양입니다. 나한텐 이 논문에서 관심 대상 우선순위가 일단 아니니.. 나중에 자세히 보도록 하겠습니다..

 

Evaluation 단계에서는 QA, Preference classification, Response generation 세 개의 task를 평가했다고 합니다. 이 각각에 해당하는 데이터셋은 시중에 나와있는 걸로 평가하지 않고 나름의 질문 구성을 했나봄.. 대단타 진짜..

 

Results

이게 결과! 되게 잘하죠~

뒤에 Ablation study 부분도 분석 잘 돼있는데 일단 여까지만 하겠음ㅋㅋ