Junction Tree Variational Autoencoder for Molecular Graph Generation

Introduction

신약개발의 핵심 과제는 원하는 화학적 특성을 가진 target 분자를 찾는 것임. 이 작업은 숙련된 화학자와 약리학자가 수년동안 개발과 탐색하는 과정이 필요함. 그러므로 이러한 과정을 automate하는 것이 궁극적인 목표임.

computational perspective에서 이 과제를 두 개의 상호 보완적인 subtasks로 분해함
1️⃣ encoding : 분자의 특성 예측과 최적화를 용이하게 하는 연속적 표현을 학습하는 것
2️⃣ decoding : 최적화된 연속 표현을 다시 향상된 특성을 가진 분자 그래프로 복원하는 것

이전 drug design 연구들은 그래프 생성 문제를 문자열 생성 문제로 변환하여 접근했음. 구체적으로 SMILES를 생성하였음.

SMILES
화학 구조를 기술하기 위해 사용되는 선형 문자열 표기법
RDKit과 같은 deterministic mapping을 이용하여 그래프로 변환할 수 있음

하지만 이러한 설계는 두 가지 한계가 있음
1️⃣ SMILES 표현은 분자 간의 유사성을 포착할 수 없음
ex ) 화학 구조가 비슷한 두 분자라도, SMILES 문자열은 다르게 인코딩 될 수 있음

이는 VAE 같은 생성 모델이 smooth molecular embedding을 학습하는 것을 방해함.

2️⃣ 중요한 화학적 속성은 선형 SMILES 표현보다 그래프 상에서 훨씬 더 명확하게 표현될 수 있음.
그러므로 그래프를 직접 다루는 것이 valid chemical structure의 생성 모델링을 더 효과적으로 만든다고 가정함

일반적으로 그래프를 node by node로 생성하는 접근법을 생각할 수 있지만 이는 분자 생성에는 적합하지 않음. 원자 단위로 분자를 하나씩 만들게 되면 chemically invalid intermediaries를 반드시 거치게 되기 때문임.

즉, 완성된 그래프가 나오기 전까지 분자의 타당성을 검증할 수없게 됨.

---

Junction Tree Variational Autoencoder

각 분자를 valid components들의 vocabulary에서 선택된 subgraphs들로 구성된 것으로 해석함. 이 components는 encoding, decoding할 때도 building block으로 사용됨. 이는 decoder가 화학적으로 타당한 분자를 유요한 구성요소들의 조합과 상호작용을 통해 부분별로 만들어낼 수 있다는 점이 핵심 장점임.

이러한 방법을 통해 항상 화학적으로 valid molecule를 유지하면서 단계별 생성이 가능하게 됨. ( 이는 앞에서나온 SMILES 한계를 해결함을 보여줌 )

Example
방향족 결합은 전체 방향족 고리가 존재하지 않는 한 그 자체로는 화학적으로 유효하지 않음. 
따라서, 방향족 고리를 원자 하나씩 만드는 대신, "고리" 자체를 기본 어휘의 일부로 도입하는 
것이 훨씬 합리적임. 

vocabulary of components
ring, bond, atom 등이 구성되며, 하나의 분자를 이러한 components 또는 clusters of atom으로 유효한 subgraph들로 나누어 표현할 수 있고, 이 subgraphs는 일부 원자들을 공유(overlap)하며 전체 분자를 완전히 표현(cover)할 수 있음

🤔 왜 분자를 junction tree로 표현해도 정보 손실이 크지 않을까?

clique(클리크)
모든 노드가 서로 직접 연결되어 있는 완전 연결

• 세 원자 A-B-C가 모두 서로 결합되어 있다면, A-B-C는 클리크 ✅
• A-B, B-C만 결합되어 있고 A-C는 아니면 클리크 ❌

cluster
subgraph를 클러스터라고 부름

• 하나의 benzene ring, amide group, carboxyl group 같은 것들이 하나의 클러스터가 됨

클러스터들은 graphical model의 클리크와 유사한 역할을 함.
➡️ 클러스터들이 충분히 표현력이 높으므로 분자는 클러스터들 간에 overlap이 있더라도 클러스터 cycle을 만들지 않고 덮일 수 있음.

[overlap 관계]
Cluster A (benzene ring) : [C1,C2,C3,C4,C5,C6]
Cluster B (methyl group) : [C6, H7,H8,H9]
overlap(공유 원자) : C6

이렇게 되면, A와 B는 C6 하나를 공유하지만, 전체 구조는 A-B 하나의 연결선으로 표현할 수 있음

[Junction Tree로 표현]
    Cluster A : benzene ring
			|
    Cluster B : methyl group

A-B는 연결되어 있지만, 전체 구조에는 사이클이 없음. 하지만 A와 B의 overlap은 존재함

이러한 의미에서, 클러스터들은 분자 그래프의 (non-optimal) triangulation에서의 클리크 역할을 수행함.

triangulation

A —— B        A —— B
|    |   	  | \  |
|    |	 ➡️	  |   \|
D —— C		  D —— C

첫 그래프는 4-cycle로 A-B-C-D-A로 고리를 이루고 있음. 이 상태로는 트리 구조로 분해가 어려움. 그래서 이 그래프를 삼각형들로만 이루어진 구조로 바꾸기 위새 A-C 엣지를 하나 추가함. 그러면 A-B-C, A-C-D 두 개의 삼각형으로 이루어지게 됨.

이렇게 triangulation을 통해 그래프가 여러 개의 클리크로 나뉘게 되고, 이 클리크들을 트리 형태로 연결하면 junction tree가 됨.

Clique1 = {A,B,C} ,  Clique2 = {A,C,D}
두 클리크는 A와 C를 overlap 하므로 다음과 같이 연결할 수 있음
[{A,B,C}] - [{A,C,D}]

이렇게 구성된 클러스터들로부터 junction tree를 형성하며, 이를 분자의 tree representation으로 사용함.

Notation

기호	설명
$G=(V,E)$	molecular graph $V$ : 원자의 집합, $E$ : 결합의 집합
$N(x)$	$x$의 이웃
$\sigma(\cdot)$	시그모이드 함수
$\tau(\cdot)$	ReLU 함수
$i,j,k$	트리 내의 노드
$u,v,w$	그래프 내의 노드

순서	내용
① Molecule $\rightarrow$ Molecular Graph G	실제 분자 구조를 그래프 $G = (V,E)$로 표현
② Molecular Graph G $\rightarrow$ $z_G$ (Graph Encoding)	Graph Encoder $q(z_G \mid G)$ ➡️ 그래프 형태의 분자를 latent vector $z_G$로 인코딩 ➡️ 분자의 세밀한 연결 관계를 학습
③ Molecular $\rightarrow$ Tree Decomposition** $\rightarrow$ Junction Tree T	분자 $\rightarrow$ Tree Decomposition 수행 ➡️ 분자를 구성하는 화학적으로 유효한 subgraph(clusters)들을 추출 ➡️ 이 클러스터들의 연결 관계를 트리 형태로 구성
④ Junction Tree T $\rightarrow$ $z_T$ (Tree Encoding)	Tree Encoder $q(z_T \mid T)$ ➡️ 트리 구조를 latent vector $z_T$로 인코딩 ➡️ "어떤 클러스터들이 존재하고, 트리가 어떤 구조를 가지는가" 정보를 저장
⑤ Latent Representation	원래의 분자 그래프(molecular graph)와 그에 대응하는 결합. 트리(junction tree)는 분자의 상호 보완적인 두 가지 표현을 제공함. 따라서 분자를 두 부분으로 이루어진 latent representation(z)으로 인코딩함 $z = [z_T,z_G]$ - $z_T$ : tree structure와 트리 내 클러스터 정보를 인코딩하지만, 클러스터들 간의 세부 연결 방식까지는 완전히 표현하지 않음 - $z_G$ : 그래프를 인코딩하여 세밀한 연결성을 포착함 그 후 latent representation은 두 단계로 molecular graph로 decoding됨.
⑥Decoding - 1단계 (Tree Decoder)	$z_T$에 포함된 정보를 바탕으로 tree decoder($p(T \mid z_T$)를 사용하여 junction tree를 재생함.
⑦ Decoding - 2단계 (Graph Decoder)	결합 트리 내의 클러스터들 사이의 세밀한 연결성을 예측하기 위해 graph decoder($p(G \mid T, z_G$))를 사용하여 full molecular graph를 구현함.

Junction Tree

Tree Decomposition이란 그래프 $G$를 cycle이 없는 트리 구조로 변환하는 것으로 원래 그래프의 여러 정점을 하나의 노드로 묶어서 전체가 트리 형태가 되도록 만드는 것임.

주어진 그래프 $G=(V,E)$에 대해, junction tree $T_G = (V_T,E_T,X)$는 다음 조건을 만족하는 연결된(labeled) 트리이며, 그 노드 집합은 $V_T = \{C_1,C_2,...,C_n\}$이고, 엣지 집합은 $E$이다.

각 노드 또는 클러스터 $C_i=(V_i,E_i)$는 그래프 $G$의 induced subgraph 이며, 아래와 같은 두가지 제약 조건들을 만족함

1️⃣ 모든 클러스터의 합집합은 $G$와 같음

2️⃣ Running intersection
모든 클러스터 $C_i,C_j,C_k$에 대하여, $C_k$가 $C_i$와 $C_j$를 잇는 경로 위에 있다면

induced subgraph들을 클러스터 레이블로 간주하므로, junction tree는 labeled tree로 표현될 수 있으며, 그 레이블의 vocabulary는 $X$이다.

molecular tree decomposition에서는 $X$에 cycle, ring, single edge만 포함되므로 vocabulary의 크기는 제한적임. (25만 개의 분자를 포함한 표준 데이터셋에는 $\mid X\mid = 789$)

Tree Decomposition of Molecules

앞에서 말했듯이 bond와 ring을 클러스터 vocabulary $X$로 사용함.

1️⃣ 주어진 분자 그래프 $G$에서 먼저 모든 simple cycle & 	고리에 속하지 않는 결합(edge)를 찾음
2️⃣ 두 개 이상의 원자를 공유하는 고리들은 하나의 bridged compound로 병합함. 
3️⃣ 이렇게 얻어진 각 고리와 결합을 하나의 클러스터로 간주 한 후, 서로 겹치는 클러스터들 사이에 엣지를 추가하여 클러스터 그래프를 구성하고, 그 중 하나의 spanning tree를 선택하여 이를 junction tree로 사용함. 

이러한 고리 병합 덕분에 junction tree 내의 임의의 두 클러스터는 최대 두 개의 원자만 공유하게 되어, 그래프 디코딩 단계에서 효율적인 추론이 가능해짐

Graph Encoder (sec 2.2)

Message Passing Network를 이용해 분자 그래프 $G$의 latent representation($z_G$)로 인코딩함.

각 정점(원자) $v$는 원자 종류, 원자가, 전자 수 등 화학적 속성을 나타내는 특징 벡터 $x_v$를 가짐.

각 엣지(결합) $(u,v) \in E$는 결합 종류(단일, 이중, 방향 등)을 나타내는 특징 벡터 $x_uv$와 메시지 방향을 구분하기 위한 hidden vector $\nu_{uv}, \nu_{vu}$를 가짐.

그래프는 일반적으로 loopy 구조를 가지므로, 노드 간의 정보를 여러 번 주고받는 loopy belief propagation 방식으로 학습함. 각 단계 $t$마다, 노드 $u \rightarrow v$로 전달되는 메시지는 다음과 같이 업데이트 됨

• $\nu_{uv}^{(t)}$ : $t$번째 반복에서 $u$가 $v$에게 보내는 방향성 메세지
• $W$ : 가중치
• $x_u$ : 정보를 보내는 쪽의 특징 정보
• $x_{uv}$ : $u$와 $v$ 사이 결합 정보
• $\sum_{w\in N(u)/v}\nu_{wu}^{(t-1)}$ : $u$가 직전 단계에서 이웃들로부터 받은 메시지들의 합
• $w\in N(u)\backslash v$ : $u$의 모든 이웃 중 $v$만 제외한 집합

$T$번의 반복 후, 이렇게 계산된 메시지들을 aggregate하여 각 정점의 latent vector를 얻음.

즉, 자기 자신의 특징 $x_u$와 이웃 노드로부터 받은 메시지들의 합을 이용해 노드의 local structural context를 표현함.

모든 노드 벡터의 평균으로 얻은 $h_G$는 분자 전체를 대표하는 graph-level embedding이 됨.

$h_G$로부터 두 개의 선형 계층(affine layers)를 통과시켜 $\mu_G$ 와 $log\ \sigma_G$를 구한 뒤 VAE 방식에 따라 가우시안 분포 $\mathcal{N}(\mu_G,\sigma_G)$에서 샘플링하여 final latent vector $z_G$를 얻음

🖊️ A-B-C인 경우가 예시라고 보자.
- t=1에서는 A는 B의 정보, B는 A,C의 정보, C는 B의 정보만 반영되어 local 정보만 반영된 상태이다

- t=2, 이전 회차에서 받은 메시지를 다시 활용해서 이웃의 이웃까지 정보를 확장할 수 있다
	A는 (B가 C로부터 받아오 정보)까지 간접적으로 포함하므로 A가 C의 영향을 받기 시작하게 된다.

즉, 이렇게 계속 반복할수록 더 먼 이웃의 정보까지 들어와 결국 노드 임베딩이 점점 더 전체 그래프 구조를 내포하게 됨.

Tree Encoder (sec 2.3)

Tree message passing network를 이용하여 인코딩함.

각 노드(클러스터) $C_i$는 클러스터의 label type을 나타내는 one-hot vector $x_i$로 표현됨.

각 엣지 ($C_i,C_j$)는 두 개의 방향성 메시지 $m_{ij},m_{ji}$를 가짐.

트리는 그래프와 달리 루프가 없으므로 두 단계 message propagation phases로 메시지를 전파함.

1️⃣ Bottom-up phase

• 임의의 리프 노드를 root로 선택
• 리프에서 root 방향으로 메시지 전파

2️⃣ Top-down phase

• 루트에서 다시 리프 방향으로 메시지 전파

이 두 단계를 통해 트리 전체에서 양방향 정보 흐름을 모두 반영할 수 있음.

각 엣지 $(i,j)$에 대한 메시지 $m_{ij}$는 다음과 같이 계산됨

즉, 노드 $i$ 자기 자신의 특징 $x_i$와 이웃 $N(i)\backslash j$로부터 들어온 메시지들을 이용하여 새로운 메시지 $m_{ij}$를 생성함.

여기서 사용된 GRU는 tree message passing용으로 변형된 GRU구조이며 아래와 같은 수식으로 정의됨.

• $s_{ij} = \sum_{k\in N(i)\backslash j}m_{kj}$
  ➡️ $i$의 자식들로부터 받은 메시지들의 합
• $z_{ij} = \sigma(W^zx_i + U^zs_{ij }+b^z)$
  ➡️ update gate - 현재 상태를 얼마나 새정보로 바꿀지?
• $r_{ki} = \sigma(W^rx_i + U^r m_{ki} + b^r)$
  ➡️ reset gate - 각 자식 메시지 $m_{ki}$를 얼마나 반영할지?
• $\tilde{m}_{ij} = tanh(Wx_i + U\sum_{k \in N(i)\backslash j}r_{ki} \odot m_{ki})$
  ➡️ 이웃 메시지들을 리셋 게이트로 조정해 요약한 후, 새로운 후보 메시지를 만듦
• $m_ij = (1-z_{ij})\odot s_{ij} + z_{ij} \odot \tilde{m}_{ij}$
  ➡️ 이전 요약($s_{ij}$)과 새로운 후보($\tilde{m}_{ij}$)를 update gate로 가중합

모든 메시지 패싱이 끝나면, 각 노드의 latent representation $h_i$는 이웃으로부터 들어온 메시지의 집합을 이용해 계산함

final tree representation은 rood node의 latent vector로 정의됨

이 벡터는 root가 지정된 트리 ($T, root$)를 인코딩함

graph encoder와 달리, tree encoder는 node average pooling을 적용하지 않음.
tree decoder가 어떤 노드를 먼저 생성해야 하는지에 대한 정보를 잃어버리기 때문임.

$z_{T_G}$는 graph encoder와 유사한 방식으로 샘플링됨. 단순화를 위해 $z_{T_G}는 z_T$로 표기함.

Tree Encoder는 프레임워크 내에서 두 가지 역할을 수행함

1️⃣ $z_T$를 계산하기 위한 용도 ➡️ bottom-up 단계만 필요(Tree Encoder)
2️⃣ 트리 $T$가 $z_T$로부터 디코딩된 후, 트리 전체에서 메시지 $m_{ij}$들을 다시 계산하여 graph decoding 동안 각 노드의 중요한 문맥 정보를 제공함. ➡️ top-down 및 bottom-up 두 단계가 모두 필요(sec 2.5)

Tree Decoder (sec 2.4)

latent representation $z_T$로부터 junction tree $T$를 decoding함. 트리는 top-down 방식으로 구성되며, 한 번에 하나의 노드를 생성해 나감

위 이미지와 같이 tree decoder는 root에서 시작하여 트리 전체를 DFS 방식으로 탐색하며 노드들을 생성함. 각 노드 방문 시 두 가지를 수행함 topological prediction, label prediction 이 후 노드가 더 이상 자식이 없다고 판단되면, backtrack 하여 이전 노드로 돌아가고, 다른 분기를 계속 확장함.

모델은 시간 $t$에 노드 $i_t$를 방문한다. $\tilde{E}_t$를 첫 $t$개의 엣지로 이루어진 부분 집합으로 정의할 때, 메시지 $h_{i_t,j_t}$는 이전 메시지들을 통해 다음과 같이 갱신됨.

• $x_{i_t}$ : 현재 노드의 특징
• $\{h_{k,i_t}\}$ : 이미 생성된 이웃 노드로부터 들어온 메시지
• $\tilde{\mathcal{E}}=\{(i_1,j_1),...,(i_m,j_m)\}$: 트리 $T=(V,E)$에 대한 DFS 동안 방문된 엣지들의 순서쌍 집합

즉, 새로운 노드를 생성할 때 이전 메시지와 현재 노드의 상태 정보를 통합하여 세밀하고 화학적으로 일관된 연결 관계를 만들어 주는 hidden state임.

Topological Prediction (위상적 예측)

모델이 노드 $i_t$를 방문할 때, 그 노드가 여전히 생성해야 할 자식이 남았는지를 binary 형태로 예측함

잠재 벡터 $z_T$, 노드의 특징 $x_{i_t}$, 그리고 현재 노드로 들어오는 메시지들 $h_{k,i_t}$를 결합하여 시그모이드 함수를 통해 예측 확률을 계산함

Label Prediction

부모 노드 $i$로부터 자식 노드 $j$가 생성될 때, 그 노드가 어떤 클러스터 라벨을 가져야할지 예측함.

$q_j$는 레이블 vocabulary $X$ 위에서의 확률 분포이다. 만약 $j$가 root node라면, 그의 부모 $i$는 가상 노드로 간주되며, 이때 $h_{ij}=0$으로 설정됨.

Learning

tree decoder는 junction tree의 정확한 구조를 복원하도록 학습되었음. 그러므로 $p(T\mid z_T)$를 최대화하는 것이 중요함.

• $\hat{p_t}$ : topological 정답값 (0 또는 1)
• $\hat{q_j}$ : label 정답값
• $L_d,L_l$ : 각각 topological prediction 과 label prediction에 대한 cross-entropy loss.

시퀀스 생성과정과 유사하게, 훈련과정에서 teacher forcing을 적용함. 각 단계에서의 topology와 label 예측 후, 모델이 이전 단계에서의 예측값 대신 ground truth을 입력으로 받아 정확한 생성 이력을 기반으로 다음 출력을 예측하도록 학습함.

Graph Decoder (sec 2.5)

모델의 마지막 단계로, 예측된 junction tree $T=(V,E)$를 기반으로 그에 대응하는 분자 그래프 $G$를 재구성하는 것임. 동일한 junction tree에 대응되는 가능한 molecule는 잠재적으로 여러 개 존재할 수 있기 때문에 deterministic하지 않음.

벤젠과 아마이드가 연결된 트리라고 해도 두 원자를 어느 위치에 결합시키느냐에 따라
실제 화학적 구조가 달라질 수 있음

그래서 트리 내의 노드(클러스터)들에 해당하는 부분 그래프들을 정확한 분자 그래프로 조립하는 것이 중요한 목표임.

그래서 Graph Decoding은 structured prediction 문제로 정의됨. 트리 $T$에 속한 모든 가능한 그래프 후보 $G'$들 중, scoring function $f^a(G')$이 가장 높은 그래프를 선택함.

• $\mathcal{G}(T)$ : junction tree $T$를 갖는 그래프들의 집합
• $f^a$ : candidate graph 들에 대한 scoring function, 이 조립 방식이 얼마나 화학적으로 자연스러운가를 평가함

scoring
하나의 트리 전체를 한꺼번에 채점하지 않고, 각 노드(클러스터) 단위로 스코어를 게산하고 합산함.

즉, 트리 $T$ 내에서 클러스터 $C_i$가 그 이웃 클러스터들 $C_j,j \in N_T(i)$에 어떻게 연결되어 있는가에만 의존함.

최고 점수를 갖는 그래프 $\hat{G}$를 찾는 것은, junction tree에 의해 유도된 graphical model 상에서의 추론 문제로 표현될 수 있음.

먼저 root와 그 이웃들의 조립을 그들의 점수에 따라 샘플링하고, 그 후 root에 연결된 이웃들과 그들의 관련된 클러스터들을 조립해 나감. ( 이때 root의 조립에 의해 설정된 degree of freedom )은 제거해야 함. 이 과정을 반복하면 됨.

트리 내의 클러스터 $C_i$와 그 이웃 클러스터들 $C_j,j \in N_T(i)$을 특정한 방식으로 결합했을 때 형성되는 subgraph를 $G_i$라고 할때, $G_i$를 candidate subgraph로 보고 점수를 부여함.

먼저 $G_i$의 vector representation $h_{G_i}$를 계산함. 이는 subgraph가 얼마나 타당한 연결인지를 평가하는 점수임.

이때 $h_{G_i}$는 message passing과 tree context 방식으로 생성이 됨.

먼저 $\alpha$인덱스로 어느 클러스터 소속 원자인지를 표시함. $u,v$를 candidate subgraph $G_i$ 안에 있는 원자들일 때

$a_v$는 두 가지에 쓰임
1. junction tree내에서 원자들의 위치를 표시하는 데 사용
2. tree encoding algorithm을 실행하여 엣지(i,j)를 따라 $i$ 아래의 subtree를 요약하는 message $\hat{m}_{i,j}$를 검색하는데 사용됨.

candidate subgraph $G_i$ 안에서도, Graph Encoder 때처럼 message passing을 통해 원자 $\cdot$ 결합 정보를 모음

위 식이 $\nu_{uv}^{(t)} = \tau(W^g_1 x_u + W^g_2x_{uv}+W^g_3\sum_{w\in N(u)\backslash v}\nu_{wu}^{(t-1)})$과의 주요 차이점은 $a_u \neq a_v$인 경우 예측된 트리 $T$에 대해 tree encoder를 실행함으로써 얻은 트리 메시지 $m_{a_u,a_v}$를 모델에 추가한다는 점임.

$\hat{m}_{a_u,a_v}$는 결합 ($u,v$)에 대한 tree-dependent positional context를 제공함.

Learning

Graph Decoder의 파라미터들은, 각 트리 노드에서 정답 그래프 $G$의 올바른 subgraph $G_i$를 예측하는 log-likelihood를 최대화하도록 학습함.

• $\mathcal{G}_i$ : 트리 노드 $i$에서 가능한 candidate subgoal 집합

Graph Decoder는 ground truth tree를 입력으로 받아 학습하여 각 트리 노드에서 여러 subgraph 후보 중 실제 결합 방식을 가장 잘 맞히도록 학습함.

Experiments

JT-VAE의 실험은 모델이 얼마나 분자를 잘 생성하고, 재구성하며, 특정 화학적 특성을 최적화할 수 있는지를 다각도로 평가하는 데 초점을 둠.

Task

• Molecule reconstruction & validity
• Bayesian optimization
• Constrained molecule optimization

Data

• ZINC 분자 데이터셋
• 약 250,000개의 약물 분자 포함

Baselines

모델	개요
Character VAE(CVAE)	SMILES 문자열을 문자 단위로 생성하는 기본 VAE
Crammar VAE(GVAE)	SMILES의 문법 제약을 문맥 자유 문법으로 반영
Syntax-directed VAE(SD-VAE)	attribute grammar를 통해 SMILES의 문법적 및 의미적 제약을 모두 통합
GraphVAE	원자 라벨과 인접 행렬을 직접 생성함

Model Configuration

latent space 차원은 56으로 통일

• 트리 인코딩($h_T$) : 28차원
• 그래프 인코딩($h_G$) : 28차원

총합 56차원 latent vector : $z = [z_T,z_G]$

Molecule Reconstruction and Validity

Setup

목적
① latent space에서 입력 분자를 다시 복원할 수 있을까?
② 아무런 조건 없이 prior 분포에서 샘플링해도 화학적으로 valid 분자가 만들어지는가?

재구성 정확도
JT-VAE는 인코딩과 디코딩이 확률적이므로 단 한 번의 인코딩/디코딩으로 정확도를 측정하기 어려움 ➡️ Monte Carlo 샘플링을 사용해 평균 정확도를 구함
1. 각 분자를 10번 인코딩 ( latent vector $z=[z_T,z_G]$ 생성 )
2. 각 인코딩 결과를 다시 10번 디코딩
3. 한 분자당 총 100개의 재구성 결과 생성
4. 이 중 입력 분자와 동일한 구조를 가진 비율을 재구성 정확도로 계산

유효성
사전 분포 $\mathcal{N}(0,1)$에서 1000개의 잠재 벡터를 무작위로 샘플링. 각 벡터를 100번 디코딩하여 분자 생성. 생성된 모든 분자를 RDKit으로 화학적 유효성 검사 & Ablation study로 디코딩 단계에서 유효성 검증을 수행하지 않은 경우의 결과도 함께 보고함.

Results

JT-VAE가 이전 모델들보다 분자 재구성 성능에서 우수함을 보여줌. 또한 prior distribution으로 부터 샘플링할 때 항상 유효한 분자(validity = 100%)를 생성함. 유효성 검증 단계를 제거하더라도, 여전히 93.5%의 유효한 분자를 생성함.

즉, 모델이 prior knowledge에 과도하게 의존하지 않음을 알 수 있음

모델이 학습 데이터를 참고하지 않고, 단순히 정규분포에서 랜덤하게 숫자를 뽑아 만든 분자들로
모두가 정상적인 유기화합물이다. 그러므로 모델이 단순히 학습 데이터를 암기하지 않고, 화학 규칙을
스스로 배워 valid하고 다양한 분자를 생성했음을 알 수 있음

Analysis

JT-VAE의 latent space를 qualitatively 분석하기 위해, 분자의 이웃 구조를 시각화하였음. 하나의 기준 분자를 선택하고 그 주변 이웃 분자들을 grid 형태로 시각화 하였음. 즉, latent space에서 가까운 위치의 분자들이 구조적으로도 비슷한지를 확인하였음.

기존 모델은 데이터에 거의 없는 7개 이상의 운자를 가진 거대 고리 구조를 생성하기도 하였지만 
JT-VAE는 그런 비현실적 구조 없이 데이터 분포에 맞는, 자연스러운 구조만 생성하였음. 또한, 
동일한 트리 구조를 유지하되 클러스터 연결 방식만 달라지는 분자 그룹들이 존재하였음.
즉, JT-VAE의 잠재공간은 smooth하며, 구조적으로 자연스러운 변화를 반영함.

Bayesian Optimization

Setup

목적
특정한 화학적 특성이 높은 새로운 분자를 찾아낼 수 있는지를 평가함.
이때 사용하는 목표 특성은 다음과 같음

$y(m) = logP-SA\ penalty - long\ cycle\ penalty$

• $logP$ : 분자의 지용성 / 친수성을 나타내는 수치
• $SA$ : 합성 난이도 ( 만들기 어려운 분자는 점수에서 패널티)
• $long\ cycle \ penalty$ : 너무 큰 고리 구조가 있는 분자는 화학적으로 불안정하므로 감점

즉, 지용성은 높고, 합성은 쉽고, 구조는 안정적인 분자를 찾아내는 것이 목표

과정
① VAE 학습시켜 각 분자를 latent vector로 인코딩하고, 이는 VAE의 평균으로 정의됨.
② VAE가 학습된 후, Sparse Gaussian Process(SGP)를 훈련시켜 latent representation을 입력받아 해당 분자의 $y(m)$을 예측하도록 함
③ 마지막으로, Expected Improvement 휴리스틱을 사용하여 batch 5 iterations에 걸쳐 베이지안 최적화를 수행

두 가지 실험 결과
1️⃣ SGP의 예측 성능

- 서로 다른 VAE 모델에서 학습된 잠재 인코딩을 이용하여 SGP를 훈련하고, log-likelihood와  RMSE로 평가함.
- 10-fold cross validation을 사용함

2️⃣ BO로 탐색한 상위 3개 분자

- 서로 다른 모델이 탐색한 결과를 비교함

Results

JT-VAE가 이전 방법들보다 훨씬 더 높은 점수를 가진 분자들을 발견했음

위 이미지는 JT-VAE가 찾은 top-3의 최고 점수 분자를 보여줌. 실제로 JT-VAE는 3.50점 이상의 점수를 가진 분자를 50개 이상 탐색했으며, 이는 SD-VAE가 제안한 두 번째로 좋은 분자보다도 훨씬 높은 성능임

JT-VAE 임베딩으로 학습된 SGP가 가장 우수한 예측 성능을 보였음.

Constrained Optimization

Setup

목적
제약조건이 있는 상황에서 분자 최적화를 수행하는 것
주어진 분자 $m$이 있을 때, 분자 유사도 $sim(m,m') \ge δ$을 만족하면서, 화학적 속성 값 $y(m')$이 가장 높은 분자 $m'$를 찾는 것

setup

• 유사도는 Tanimoto similarity를 사용
• Morgan fingerprint로 계산됨
• target property는 penalized logP 계수를 사용함.
• JT-VAE와 함께 property predictor F를 공동학습 함. $F$는 feed forward neural network로 매개변수화되어 있으며, 각 분자의 latent embedding로 부터 해당 분자의 속성 $y(m)$을 예측함

Optimization
1️⃣ 주어진 분자 $m$의 latent representation에서 시작함
2️⃣ 예측된 점수 $F(\cdot)$을 높이기 위해 latent space에서 gradient ascent를 수행함.
3️⃣ $K = 80$번의 gradient step을 적용한 후, 각 단계에서 얻어진 latent vector들을 디코딩하여 k개의 candidate 분자를 얻음
4️⃣ 이 중에서, 유사도 제약을 만족하면서 가장 높은 $F(\cdot)$값을 가진 분자를 최종 선택

기준
modification이 성공했다고 판단하는 기준은 다음과 같음
디코딩된 분자들 중 하나 가
(a) 유사도 제약 $sim(m,m') \ge δ$를 만족하고,
(b) 원래 분자 $m$과 서로 다른 분자인 경우

실험에서는 가장 어려운 상황을 설정하기 위해, 테스트 세트에서 가장 낮은 특성 점수 $y(\cdot)$을 가진 800개의 분자를 선택하였음. 다음 세 가지 지표를 report함.

• success rate : 수정이 성공한 비율
• average improvement : $y(m') - y(m)$
• molecular similarity : $sim(m,m')$

Results

• 비제약 상황 ( $\delta = 0$ )
  average improvement가 가장 높았으나, 생성된 분자들이 원래 분자와 매우 다른 경우가 많았음
  
• 제약을 강화 ( $\delta = 0.4$ )
  모델은 약 80%의 경우에서 원래 분자와 유사한 분자를 찾아냈으며, 평균 향상도는 0.84였음
  

이는 모델이 학습한 latent space가 smooth하고 연속적인 구조를 가지고 있음을 보여줌

JT-VAE가 유사도는 유지하면서도 특성 향상이 큰 효과적인 분자 수정 사례를 시각적으로 보여줌.

Conclusion

JT-VAE는 분자 생성과 분자 최적화 두 측면 모두에서 기존 연구들을 유의미하게 능가하였으며 향후 연구에서는 이 방법을 보다 일반적인 low-treewidth로 일반화하는 것을 목표로 함.