https://github.com/srush/Transformer-Puzzles

Puzzles for exploring transformers
https://github.com/srush/Transformer-Puzzles

Last synced: 8 days ago
JSON representation

Puzzles for exploring transformers

• Host: GitHub
• URL: https://github.com/srush/Transformer-Puzzles
• Owner: srush
• License: mit
• Created: 2023-05-01T07:09:28.000Z (over 1 year ago)
• Default Branch: main
• Last Pushed: 2023-05-04T18:55:23.000Z (over 1 year ago)
• Last Synced: 2024-10-15T09:34:13.277Z (28 days ago)
• Language: Jupyter Notebook
• Size: 147 KB
• Stars: 315
• Watchers: 7
• Forks: 23
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

        
# Transformer Puzzles



  



This notebook is a collection of short coding puzzles based on the internals of the Transformer. The puzzles are written in Python and can be done in this notebook. After completing these you will have a much better intutive sense of how a Transformer can compute certain logical operations. 

These puzzles are based on [Thinking Like Transformers](https://arxiv.org/pdf/2106.06981.pdf) by Gail Weiss, Yoav Goldberg, Eran Yahav and derived from this [blog post](https://srush.github.io/raspy/).

![image](https://user-images.githubusercontent.com/35882/235678934-44c83052-9743-4de7-a46c-49a517923da1.png)

## Goal

**Can we produce a Transformer that does basic elementary school addition?**

i.e. given a string "19492+23919" can we produce the correct output? 

## Rules

Each exercise consists of a function with a argument `seq` and output `seq`. Like a transformer we cannot change length. Operations need to act on the entire sequence in parallel. There is a global `indices` which tells use the position in the sequence. If we want to do something different on certain positions we can use `where` like in Numpy or PyTorch. To run the seq we need to give it an initial input. 

```python colab={"base_uri": "https://localhost:8080/", "height": 96} id="1b28dc98" outputId="f1ac1157-3db8-40c0-dbb2-7d9bad8943a0"

def even_vals(seq=tokens):

    "Keep even positions, set odd positions to -1"

    x = indices % 2

    # Note that all operations broadcast so you can use scalars.

    return where(x == 0, seq, -1)

seq = even_vals()

# Give the initial input tokens

seq.input([0,1,2,3,4])

```

The main operation you can use is "attention". You do this by defining a selector which forms a matrix based on `key` and `query`.

```python colab={"base_uri": "https://localhost:8080/", "height": 176} id="e2ee0ff8" outputId="a61ac19c-2550-4f3c-d653-50c323cdfd59"

before = key(indices) < query(indices)

before

```

We can combine selectors with logical operations.

```python colab={"base_uri": "https://localhost:8080/", "height": 201} id="c315ba6d" outputId="270d50fa-649c-438b-8606-d3d078478162"

before_or_same = before | (key(indices) == query(indices))

before_or_same

```

Once you have a selector, you can apply "attention" to sum over the grey positions. For example to compute cumulative such we run the following function. 

```python colab={"base_uri": "https://localhost:8080/", "height": 326} id="e79c8c8b" outputId="44db7f90-502d-497c-c5ba-4062c09f0a9a"

def cumsum(seq=tokens):

    return before_or_same.value(seq)

seq = cumsum()

seq.input([0, 1, 2, 3, 4])

```

Good luck!