Skip to main content

TinyModel

TinyModel is a 4 layer, 44M parameter model trained on TinyStories V2 for mechanistic interpretability. It uses ReLU activations and no layernorms. It comes with trained SAEs and transcoders.

It can be installed with pip install tinymodel

from tiny_model import TinyModel, tokenizer

lm = TinyModel()

# for inference
tok_ids, attn_mask = tokenizer(['Once upon a time', 'In the forest'])
logprobs = lm(tok_ids)

# Get SAE/transcoder acts
# See 'SAEs/Transcoders' section for more information.
feature_acts = lm['M1N123'](tok_ids)
all_feat_acts = lm['M2'](tok_ids)

# Generation
lm.generate('Once upon a time, Ada was happily walking through a magical forest with')

# To decode tok_ids you can use
tokenizer.decode(tok_ids)

It was trained for 3 epochs on a preprocessed version of TinyStoriesV2. Pre-tokenized dataset here. I recommend using this dataset for getting SAE/transcoder activations.

SAEs/transcoders

Some sparse SAEs/transcoders are provided along with the model.

For example, acts = lm['M2N100'](tok_ids)

To get sparse acts, choose which part of the transformer block you want to look at (currently sparse MLP/transcoder and SAEs on attention out are available, under the tags 'M' and 'A' respectively). Residual stream and MLP out SAEs exist, they just haven't been added yet, bug me on e.g. Twitter if you want this to happen fast.

Then, add the layer. A sparse MLP at layer 2 would be 'M2'. Finally, optionally add a particular neuron. For example 'M0N10000'.

Tokenization

Tokenization is done as follows:

  • the top-10K most frequent tokens using the GPT-NeoX tokenizer are selected and sorted by frequency.
  • To tokenize a document, first tokenize with the GPT-NeoX tokenizer. Then replace tokens not in the top 10K tokens with a special [UNK] token id. All token ids are then mapped to be between 1 and 10K, roughly sorted from most frequent to least.
  • Finally, prepend the document with a [BEGIN] token id.

Release files for tinymodel 0.1.1.post13

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for tinymodel 0.1.1.post13
File Size Uploaded
tinymodel-0.1.1.post13.tar.gz 78.5 kB Details

Built distribution (wheel)

Table of built distributions (wheels) for tinymodel 0.1.1.post13
File Interpreter ABI Platform
tinymodel-0.1.1.post13-py3-none-any.whl Python 3 none any Details

Total release size: 155.9 kB

Release files / tinymodel-0.1.1.post13.tar.gz

Download URL tinymodel-0.1.1.post13.tar.gz
Size 78.5 kB
Tags Source
SHA-256 checksum
How to use checksums
94299ba8055b4ecc4fa7667fd289ad0268f958703aa3a5201d73e0e44faa88c5
BLAKE2b-256 checksum
How to use checksums
0e88b1aa9b3d8bc4189a663aee8df96ce00e734cec102490dae6b6adb0361b19
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via poetry/1.8.3 CPython/3.10.12 Linux/6.5.0-41-generic

Release files / tinymodel-0.1.1.post13-py3-none-any.whl

Download URL tinymodel-0.1.1.post13-py3-none-any.whl
Size 77.4 kB
Tags Python 3
SHA-256 checksum
How to use checksums
ce16c760ea104af47b383bf0ea996fa4aa79226ce78558fa3ac9dd169123f053
BLAKE2b-256 checksum
How to use checksums
f066f828c4be83cd01f0508b9810bef4a13babb8a0440fda39eb59f9c9edc9c5
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via poetry/1.8.3 CPython/3.10.12 Linux/6.5.0-41-generic
Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page