LLMを使って論文のMindmapを作りたい ver.1

今回は最近やっているLLMを使って論文からMindmapを作る取り組みについて紹介したいと思います。まだ試作段階で一部うまくいって、一部うまくいってないという感じですが、参考になる人もいるかもしれないと思ったので記事にまとめました。

今回のコードは以下のところに置いてあります。コードが気になる方はこちらをご覧ください

https://github.com/shu65/langchain_examples/blob/main/Make_Mind_Map_with_LangChain_v1.ipynb

1 Mindmapとは？
2 LLMからMindmapを作る取り組み
3 どうやって論文のMindmapを作るか？
4 現状の方法の何がまずいのか？
5 終わりに

Mindmapとは？

Mindmapとしてはmain topicからアイディアや思考を分岐させて可視化する手法になります。様々な場所で利用されている方法で、私もちゃんと理解した本の場合は、読んだまとめをMindmapにするということをよくやります。

そんなMindmapですが、文章の内容を理解するのに結構役立つと思っていて、もしかしたら論文のMindmapをLLMに作ってもらってそれを読めば、論文の理解スピードがあがるのではということを思っています。今回はこの仮説を確かめてみる一環で、LLMでMindmapを作らせてみよう！ということでやってみました。

LLMからMindmapを作る取り組み

LLMでMindmapを作る例はいくつかあります。今回はその中でもこれを参考にしました。

https://xmind.app/blog/chatgpt-and-xmind-how-to-create-a-mind-map-with-chatgpt/

これはXmindというMindmapツールのblog記事で、ChatGPTを使ってトピックに沿ったMindmapをChatGPTに作らせてそれをXmindに取り込んで修正するまでの流れが示されています。

使ったpromptの例もちゃんと書いてあって参考にしやすそうだったので、この記事の例を利用して今回は試作しました。

今回作ったものを次に説明していきます。

どうやって論文のMindmapを作るか？

いろいろ工夫できる部分はあると思いますが、今回はひとまず一番手っ取り早くやるならどうやるか？というのを考えて実装したものを紹介します。

手順としては以下の手順で作ることにします。

PDFからページごとにテキストを抽出
ページの量が大きすぎる場合LLMの入力できるトークン数の上限を超える可能性があるので、分割する
分割されたテキスト毎にMindmapを生成
できたMindmapをLLMにマージさせて最終的なMindmapを作る

今回は例としてちょうど最近読んだ以下の論文で試してみた例を示します。

Dao, T., Fu, D.Y., Ermon, S., Rudra, A., & R’e, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. ArXiv, abs/2205.14135.

また実装に関してはLangChainを使ってやりました。LangChainについてご存じない方は以前書いた紹介記事もありますので興味があれば見てください。

ChatGPTみたいなアプリが簡単に作れるLangChainがすごすぎ

次に各ステップに関して簡単に説明していきます。

PDFからページごとにテキストを抽出とページの分割

LangChainにはPDFを読み込む機能とページを分割する機能があるので、今回はこれをそのまま使います。コード的には以下の通りです。

from langchain.document_loaders.pdf import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

file_path = "./reserch_paper.pdf"
loader = PyPDFLoader(file_path)
pages = loader.load()
pages = pages[:11]
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    encoding_name="p50k_base",
    chunk_size=1000,
    chunk_overlap=100,
)
sub_pages = text_splitter.split_documents(pages)

まず、PyPDFLoaderで指定されたPDFを読み込みます。その後、Referenceよりも手前のページまでを取り出します。今回は11ページ目からReferenceが始まっているので、その手前のページまで取り出します。

その後、RecursiveCharacterTextSplitterでテキストを分割します。

分割されたテキスト毎にMindmapを生成

分割されたテキストごとにMindmapを作らせます。ここではXmindのblogを参考にして以下のようなpromptのテンプレートを利用しました。

MAKE_MIND_MAP_PROMPT_TEMPLATE_STR = """I would like to create a mind map using the Xmind tool about following text.
Can you provide me with some text in Markdown format that is compatible with Xmind? 
Please include a Central Topic with Main Topics and any additional information goes to Subtopics that will help create an effective mind map."

The mindmap format is follow:

# Central Tpoic Title

## Main Topic Title 1

### Subtopic Title 1

- Subtopic Title 1

  - Subtopic Title 1
  - Subtopic Title 2

### Subtopic Title 2

- Subtopic Title 1

  - Subtopic Title 1
  - Subtopic Title 2

## Main Topic Title 2

- Subtopic Title 1

  - Subtopic Title 1
  - Subtopic Title 2


Text:


{text}


Mind map in Markdown format:"""

これをLangChainのLLMChainを使ってLLMに渡します。

これでうまくいくの？と疑問に思う方もいると思うので、最初の分割（AbstractとIntroductionの最初の部分）の結果を以下に示します。

# FlashAttention : Fast and Memory-Efficient Exact Attention with IO-Awareness

## Introduction

### FLOP Reduction

- Sparse-Approximation
- Low-Rank Approximation
- Combinations

### IO-Awareness

- Reads and Writes
- Fast and Slow Memory
- GPU On-Chip SRAM and HBM

## Proposed Method

### FlashAttention

- Tiling
- Reduce Memory Reads/Writes
- Analyze IO Complexity

### Block-Sparse FlashAttention

- Faster than Existing Approximate Attention
- End-to-End Wall-Clock Speedup
- Longer Context in Transformers
- Higher Quality Models
- New Capabilities

可もなく不可もなくという感じのMindmapができているかなと思います。ちなみにXmindでimportできるフォーマットで出力されているので、画像にすることも可能です。

MindmapをLLMにマージさせて最終的なMindmapを作る

ここまでで各ページのMindmapができました。ここから1つのMindmapにまとめていく作業をさせます。

MindmapをLLMにマージしてもらうために、今回はLLMに入力できるトークン数までMindmapを連結して、以下のpromptのテンプレートに流し込んでLLMに入力しました。

"""Compact the following mind map into one mind maps. 
The format should be Xmind compatible Markdwon format.
The mind map format is follow:

# Central Tpoic Title

## Main Topic Title 1

### Subtopic Title 1

- Subtopic Title 1

  - Subtopic Title 1
  - Subtopic Title 2

### Subtopic Title 2

- Subtopic Title 1

  - Subtopic Title 1
  - Subtopic Title 2

## Main Topic Title 2

- Subtopic Title 1

  - Subtopic Title 1
  - Subtopic Title 2


{text}


Output mind map in Markdown format:"""

今回の例では全部同時に1つにマージしようとするとトークン数の限界を超えるので、Mindmapをマージしたあと、できたもの同士をまたマージするという多段階のマージを繰り返し、最終的な結果をえるようにしました。ただ、これは現状ではあまりうまくいってません。結果としては以下の通りです。

# FlashAttention : Fast and Memory-Efficient Exact Attention with IO-Awareness

## Introduction

### FLOP Reduction

- Sparse-Approximation
- Low-Rank Approximation
- Combinations
- IO-Awareness
  - Reads and Writes
  - Fast and Slow Memory
  - GPU On-Chip SRAM and HBM

## Proposed Method

### FlashAttention

- Tiling
- Reduce Memory Reads/Writes
- Analyze IO Complexity
- Bandwidth & Memory Size
- Attention on GPT-2
  - FlashAttention PyTorch
    - Time (ms)
    - Matmul
    - Mask
    - Softmax
    - Dropout
    - Matmul
    - Fused Kernel
    - Q: N x d
    - V: N X d
    - KT: d x N
    - QKT: N x N
    - sm(Q KT)V: N x d
- Outer Loop
  - Copy Block to SRAM
  - Copy
    - Outer Loop
    - Inner Loop
  - Compute Block on SRAM
  - Output to HBM
- Inner Loop
  - Outer Loop
  - GPU
    - SRAM
    - GPU
    - HBM
    - Main Memory (CPU DRAM)
  - SRAM : 19 TB/s (20 MB)
  - HBM: 1.5 TB/s (40 GB)
  - DRAM : 12.8 GB/s (>1 TB)
- Fewer HBM Accesses
  - Fig. 2
  - Lower Bound
- Memory Access Overhead
  - Block-Sparse FlashAttention
    - 2-4x Faster
    - Sequence Length of 64k
    - IO Complexity Better than FlashAttention
- Softmax Decomposition
  - Vectors 𝑥¹1º–𝑥¹2º2R𝐵
    - Softmax of Concatenated 𝑥
      - Decomposition
        - 𝑚¹𝑥º
        - 𝑓¹𝑥º
        - ℓ¹𝑥º
- Text

ぱっと見ただけで後半が変なことになっていることが分かります。これに関しては次で考察します。

現状の方法の何がまずいのか？

うまくいっていない原因としては以下のものがあります。

PDFからテキスト抽出する時点で、数式や図、テーブルが含まれる部分はそもそもテキストが変
Mindmapが大きいと出力できるトークン数の制限にひっかかるため変なところで止まる

１に関してはPDFを使っている限り常について回る問題だと思いますが、PDFのパースはちゃんとやろうとするとかなり大変です。この結果、どうしても一部変なテキストになります。LLMならもしかしたらそれでもちゃんとMindmapにできるんじゃないか？と思ってためしましたが、やはり、変なテキストの部分で作ったMindmapはクオリティがかなり低い印象です。これが最後まで響いて最終的なMindmapが変なことになっていると考えられます。

ただ、この部分に関してはpromptを工夫すればもう少しましになるかもしれないので、また今度改良しようと思っています。

また、もう一つ変なところで止まっている原因として出力できるトークン数の限界の問題があります。お金の関係もあり現在は512トークンまでで制限しています。この結果、変なところで出力がとまってしまうという問題がでています。解決方法としてはもっとたくさんトークンを出力できるようにすればいいだけではありまがすが、その場合お金がすごいことになります。そうなってくると気軽にMindmapが作れなくなってしまって微妙だなぁという感じがするので、この方法以外にいい方法がないか現在模索中です。