self-reasoning-tokens-pytorch - GitLite

KO | EN

gitlite — search

#javascript #python #react #typescript #golang #android #machine-learning #macos #windows #awesome #cli #kubernetes #awesome-list #mcp

self-reasoning-tokens-pytorch

(★ 57)

Exploration into the proposed "Self Reasoning Tokens" by Felipe Bonetto

파일 탐색기

최종 버전 다운로드 (.zip)

- - python-publish.yml
- __init__.py
- attention_with_stop_graddable_qkv.py
- self_reasoning_tokens.py
.gitignore
LICENSE
mask.png
README.md
reasoning-tokens.png
setup.py

# 설치 가이드

1. 코드 내려받기

git clone https://github.com/lucidrains/self-reasoning-tokens-pytorch

깃허브에서 프로젝트 코드 전체를 내 컴퓨터로 내려받습니다.

cd self-reasoning-tokens-pytorch

방금 내려받은 프로젝트 폴더 안으로 이동합니다.

2. Python

쉬움 추천

사전 준비물

Git GitHub에서 프로젝트 코드를 내려받으려면 필요합니다.
Python 3 설치 시 'Add Python to PATH' 옵션을 꼭 체크하세요 (Windows).

$ pip install self-reasoning-tokens-pytorch

requirements.txt 등에 명시된 파이썬 라이브러리를 설치합니다.

✅ 에러 메시지 없이 실행되고 터미널에 안내 문구가 출력되면 정상입니다.

이 레포의 README에 적힌 실제 명령어를 그대로 가져왔습니다.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<img src="./reasoning-tokens.png" width="400px"></img>

## Self Reasoning Tokens - Pytorch (wip)

Exploration into the proposed <a href="https://reasoning-tokens.ghost.io/reasoning-tokens/">Self Reasoning Tokens</a> by Felipe Bonetto. The blog post seems a bit unfleshed out, but the idea of stop gradients from next token(s) is an interesting one.

My initial thought was to apply a stop gradient mask on the attention matrix, but then realized that the values of the "reasoning" tokens could not be stop gradiented correctly without memory issues.

While walking the dog and meditating on this, I came to the realization that one can create independent stop gradient masks for queries, keys, values in either flash attention or a custom attention backwards, and there may be a whole array of possibilities there. If any experiments come back positive from this exploration, will build out a concrete implementation of this.

## Install

```bash
$ pip install self-reasoning-tokens-pytorch
```

## Usage

```python
import torch
from self_reasoning_tokens_pytorch import Transformer

model = Transformer(
    dim = 512,
    depth = 4,
    num_tokens = 256,
    stop_grad_next_tokens_to_reason = True
)

x = torch.randint(0, 256, (1, 4))

loss = model(
    x,
    num_reason_tokens = 4,                # number of reasoning tokens per time step
    num_steps_future_can_use_reason = 16, # say you wish for reason tokens to be only attended to by tokens 16 time steps into the future
    return_loss = True
)

loss.backward()

logits = model(x, num_reason_tokens = 4)
```

Or use the novel attention with ability to pass specific stop gradient masks for queries, keys, values

```python
import torch
from self_reasoning_tokens_pytorch import stop_graddable_attn

q = torch.randn(2, 8, 1024, 64)
k = torch.randn(2, 8, 1024, 64)
v = torch.randn(2, 8, 1024, 64)

stop_grad_mask = torch.randint(0, 2, (8, 1024, 1024)).bool()

out = stop_graddable_attn(
    q, k, v, causal = True,
    q_stop_grad_mask = stop_grad_mask,
    k_stop_grad_mask = stop_grad_mask,
    v_stop_grad_mask = stop_grad_mask
)

out.shape # (2, 8, 1024, 64)
```

The mask should look something like

<img src="./mask.png" width="200px"></img>

## Todo

- [ ] deviating from blog post, also try optimizing only a subset of attention heads by tokens far into the future

## Citations

```bibtex
@misc{Bonetto2024,
    author  = {Felipe Bonetto},
    url     = {https://reasoning-tokens.ghost.io/reasoning-tokens/}
}
```