declarative-attention 0.0.4__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- declarative_attention-0.0.4/.gitignore +220 -0
- declarative_attention-0.0.4/LICENSE +21 -0
- declarative_attention-0.0.4/PKG-INFO +184 -0
- declarative_attention-0.0.4/README.md +139 -0
- declarative_attention-0.0.4/declarative_attention/__init__.py +7 -0
- declarative_attention-0.0.4/declarative_attention/declarative_attention.py +198 -0
- declarative_attention-0.0.4/pyproject.toml +67 -0
|
@@ -0,0 +1,220 @@
|
|
|
1
|
+
*.sh
|
|
2
|
+
|
|
3
|
+
# Byte-compiled / optimized / DLL files
|
|
4
|
+
__pycache__/
|
|
5
|
+
*.py[codz]
|
|
6
|
+
*$py.class
|
|
7
|
+
|
|
8
|
+
# C extensions
|
|
9
|
+
*.so
|
|
10
|
+
|
|
11
|
+
# Distribution / packaging
|
|
12
|
+
.Python
|
|
13
|
+
build/
|
|
14
|
+
develop-eggs/
|
|
15
|
+
dist/
|
|
16
|
+
downloads/
|
|
17
|
+
eggs/
|
|
18
|
+
.eggs/
|
|
19
|
+
lib/
|
|
20
|
+
lib64/
|
|
21
|
+
parts/
|
|
22
|
+
sdist/
|
|
23
|
+
var/
|
|
24
|
+
wheels/
|
|
25
|
+
share/python-wheels/
|
|
26
|
+
*.egg-info/
|
|
27
|
+
.installed.cfg
|
|
28
|
+
*.egg
|
|
29
|
+
MANIFEST
|
|
30
|
+
|
|
31
|
+
# PyInstaller
|
|
32
|
+
# Usually these files are written by a python script from a template
|
|
33
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
34
|
+
*.manifest
|
|
35
|
+
*.spec
|
|
36
|
+
|
|
37
|
+
# Installer logs
|
|
38
|
+
pip-log.txt
|
|
39
|
+
pip-delete-this-directory.txt
|
|
40
|
+
|
|
41
|
+
# Unit test / coverage reports
|
|
42
|
+
htmlcov/
|
|
43
|
+
.tox/
|
|
44
|
+
.nox/
|
|
45
|
+
.coverage
|
|
46
|
+
.coverage.*
|
|
47
|
+
.cache
|
|
48
|
+
nosetests.xml
|
|
49
|
+
coverage.xml
|
|
50
|
+
*.cover
|
|
51
|
+
*.py.cover
|
|
52
|
+
.hypothesis/
|
|
53
|
+
.pytest_cache/
|
|
54
|
+
cover/
|
|
55
|
+
|
|
56
|
+
# Translations
|
|
57
|
+
*.mo
|
|
58
|
+
*.pot
|
|
59
|
+
|
|
60
|
+
# Django stuff:
|
|
61
|
+
*.log
|
|
62
|
+
local_settings.py
|
|
63
|
+
db.sqlite3
|
|
64
|
+
db.sqlite3-journal
|
|
65
|
+
|
|
66
|
+
# Flask stuff:
|
|
67
|
+
instance/
|
|
68
|
+
.webassets-cache
|
|
69
|
+
|
|
70
|
+
# Scrapy stuff:
|
|
71
|
+
.scrapy
|
|
72
|
+
|
|
73
|
+
# Sphinx documentation
|
|
74
|
+
docs/_build/
|
|
75
|
+
|
|
76
|
+
# PyBuilder
|
|
77
|
+
.pybuilder/
|
|
78
|
+
target/
|
|
79
|
+
|
|
80
|
+
# Jupyter Notebook
|
|
81
|
+
.ipynb_checkpoints
|
|
82
|
+
|
|
83
|
+
# IPython
|
|
84
|
+
profile_default/
|
|
85
|
+
ipython_config.py
|
|
86
|
+
|
|
87
|
+
# pyenv
|
|
88
|
+
# For a library or package, you might want to ignore these files since the code is
|
|
89
|
+
# intended to run in multiple environments; otherwise, check them in:
|
|
90
|
+
# .python-version
|
|
91
|
+
|
|
92
|
+
# pipenv
|
|
93
|
+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
|
94
|
+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
|
95
|
+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
|
96
|
+
# install all needed dependencies.
|
|
97
|
+
# Pipfile.lock
|
|
98
|
+
|
|
99
|
+
# UV
|
|
100
|
+
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
|
101
|
+
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
|
102
|
+
# commonly ignored for libraries.
|
|
103
|
+
# uv.lock
|
|
104
|
+
|
|
105
|
+
# poetry
|
|
106
|
+
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
|
107
|
+
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
|
108
|
+
# commonly ignored for libraries.
|
|
109
|
+
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
|
110
|
+
# poetry.lock
|
|
111
|
+
# poetry.toml
|
|
112
|
+
|
|
113
|
+
# pdm
|
|
114
|
+
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
|
115
|
+
# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
|
|
116
|
+
# https://pdm-project.org/en/latest/usage/project/#working-with-version-control
|
|
117
|
+
# pdm.lock
|
|
118
|
+
# pdm.toml
|
|
119
|
+
.pdm-python
|
|
120
|
+
.pdm-build/
|
|
121
|
+
|
|
122
|
+
# pixi
|
|
123
|
+
# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
|
|
124
|
+
# pixi.lock
|
|
125
|
+
# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
|
|
126
|
+
# in the .venv directory. It is recommended not to include this directory in version control.
|
|
127
|
+
.pixi
|
|
128
|
+
|
|
129
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
|
130
|
+
__pypackages__/
|
|
131
|
+
|
|
132
|
+
# Celery stuff
|
|
133
|
+
celerybeat-schedule
|
|
134
|
+
celerybeat.pid
|
|
135
|
+
|
|
136
|
+
# Redis
|
|
137
|
+
*.rdb
|
|
138
|
+
*.aof
|
|
139
|
+
*.pid
|
|
140
|
+
|
|
141
|
+
# RabbitMQ
|
|
142
|
+
mnesia/
|
|
143
|
+
rabbitmq/
|
|
144
|
+
rabbitmq-data/
|
|
145
|
+
|
|
146
|
+
# ActiveMQ
|
|
147
|
+
activemq-data/
|
|
148
|
+
|
|
149
|
+
# SageMath parsed files
|
|
150
|
+
*.sage.py
|
|
151
|
+
|
|
152
|
+
# Environments
|
|
153
|
+
.env
|
|
154
|
+
.envrc
|
|
155
|
+
.venv
|
|
156
|
+
env/
|
|
157
|
+
venv/
|
|
158
|
+
ENV/
|
|
159
|
+
env.bak/
|
|
160
|
+
venv.bak/
|
|
161
|
+
|
|
162
|
+
# Spyder project settings
|
|
163
|
+
.spyderproject
|
|
164
|
+
.spyproject
|
|
165
|
+
|
|
166
|
+
# Rope project settings
|
|
167
|
+
.ropeproject
|
|
168
|
+
|
|
169
|
+
# mkdocs documentation
|
|
170
|
+
/site
|
|
171
|
+
|
|
172
|
+
# mypy
|
|
173
|
+
.mypy_cache/
|
|
174
|
+
.dmypy.json
|
|
175
|
+
dmypy.json
|
|
176
|
+
|
|
177
|
+
# Pyre type checker
|
|
178
|
+
.pyre/
|
|
179
|
+
|
|
180
|
+
# pytype static type analyzer
|
|
181
|
+
.pytype/
|
|
182
|
+
|
|
183
|
+
# Cython debug symbols
|
|
184
|
+
cython_debug/
|
|
185
|
+
|
|
186
|
+
# PyCharm
|
|
187
|
+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
|
188
|
+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
|
189
|
+
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
|
190
|
+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
|
191
|
+
# .idea/
|
|
192
|
+
|
|
193
|
+
# Abstra
|
|
194
|
+
# Abstra is an AI-powered process automation framework.
|
|
195
|
+
# Ignore directories containing user credentials, local state, and settings.
|
|
196
|
+
# Learn more at https://abstra.io/docs
|
|
197
|
+
.abstra/
|
|
198
|
+
|
|
199
|
+
# Visual Studio Code
|
|
200
|
+
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
|
|
201
|
+
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
|
202
|
+
# and can be added to the global gitignore or merged into this file. However, if you prefer,
|
|
203
|
+
# you could uncomment the following to ignore the entire vscode folder
|
|
204
|
+
# .vscode/
|
|
205
|
+
# Temporary file for partial code execution
|
|
206
|
+
tempCodeRunnerFile.py
|
|
207
|
+
|
|
208
|
+
# Ruff stuff:
|
|
209
|
+
.ruff_cache/
|
|
210
|
+
|
|
211
|
+
# PyPI configuration file
|
|
212
|
+
.pypirc
|
|
213
|
+
|
|
214
|
+
# Marimo
|
|
215
|
+
marimo/_static/
|
|
216
|
+
marimo/_lsp/
|
|
217
|
+
__marimo__/
|
|
218
|
+
|
|
219
|
+
# Streamlit
|
|
220
|
+
.streamlit/secrets.toml
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Phil Wang
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,184 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: declarative-attention
|
|
3
|
+
Version: 0.0.4
|
|
4
|
+
Summary: Declarative Attention - Pytorch
|
|
5
|
+
Project-URL: Homepage, https://pypi.org/project/declarative-attention/
|
|
6
|
+
Project-URL: Repository, https://github.com/lucidrains/declarative-attention
|
|
7
|
+
Author-email: Phil Wang <lucidrains@gmail.com>
|
|
8
|
+
License: MIT License
|
|
9
|
+
|
|
10
|
+
Copyright (c) 2026 Phil Wang
|
|
11
|
+
|
|
12
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
13
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
14
|
+
in the Software without restriction, including without limitation the rights
|
|
15
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
16
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
17
|
+
furnished to do so, subject to the following conditions:
|
|
18
|
+
|
|
19
|
+
The above copyright notice and this permission notice shall be included in all
|
|
20
|
+
copies or substantial portions of the Software.
|
|
21
|
+
|
|
22
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
23
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
24
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
25
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
26
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
27
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
28
|
+
SOFTWARE.
|
|
29
|
+
License-File: LICENSE
|
|
30
|
+
Keywords: artificial intelligence,attention mechanisms,deep learning
|
|
31
|
+
Classifier: Development Status :: 4 - Beta
|
|
32
|
+
Classifier: Intended Audience :: Developers
|
|
33
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
34
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
35
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
36
|
+
Requires-Python: >=3.10
|
|
37
|
+
Requires-Dist: einops>=0.8.1
|
|
38
|
+
Requires-Dist: python-statemachine>=2.5
|
|
39
|
+
Requires-Dist: torch-einops-utils
|
|
40
|
+
Requires-Dist: torch>=2.5
|
|
41
|
+
Provides-Extra: examples
|
|
42
|
+
Provides-Extra: test
|
|
43
|
+
Requires-Dist: pytest; extra == 'test'
|
|
44
|
+
Description-Content-Type: text/markdown
|
|
45
|
+
|
|
46
|
+
<img src="./fig1.png" width="400px"></img>
|
|
47
|
+
|
|
48
|
+
## Declarative Attention - (wip)
|
|
49
|
+
|
|
50
|
+
Implementation of the procedure in [Language Models Can Control Their Own Attention](https://arxiv.org/abs/2609.02737), from Namgyu Ho et al. of KAIST AI
|
|
51
|
+
|
|
52
|
+
## Install
|
|
53
|
+
|
|
54
|
+
```bash
|
|
55
|
+
$ pip install declarative-attention
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
## Usage
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
import torch
|
|
62
|
+
from declarative_attention import DeclarativeStateMachine
|
|
63
|
+
|
|
64
|
+
# Define context chunk spans (1-indexed)
|
|
65
|
+
chunk_spans = [
|
|
66
|
+
(16, 1024), # Chunk 1
|
|
67
|
+
(1024, 2048), # Chunk 2
|
|
68
|
+
(2048, 3072), # Chunk 3
|
|
69
|
+
]
|
|
70
|
+
|
|
71
|
+
state_machine = DeclarativeStateMachine(
|
|
72
|
+
chunk_spans = chunk_spans,
|
|
73
|
+
tokenizer_decode = tokenizer.decode # or lambda t: tokenizer.decode([t])
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
# Decode loop: feed generated tokens (int, torch.long scalar, or str)
|
|
77
|
+
state_machine.step('<focus chunks="1,2">')
|
|
78
|
+
|
|
79
|
+
assert state_machine.is_focus
|
|
80
|
+
assert state_machine.active_chunks == {1, 2}
|
|
81
|
+
|
|
82
|
+
# Obtain 1D boolean attention mask for current decode step
|
|
83
|
+
mask = state_machine.get_mask(total_len = 4096) # True for kept keys, False for masked out
|
|
84
|
+
|
|
85
|
+
# Model finishes chunk reasoning and reverts back to global
|
|
86
|
+
state_machine.step('</focus>')
|
|
87
|
+
assert state_machine.is_global
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
### Custom Attention Patterns
|
|
91
|
+
|
|
92
|
+
Researchers can easily declare custom attention patterns and tags.
|
|
93
|
+
|
|
94
|
+
#### 1. Logarithmically Spaced Attention (Recent to Past)
|
|
95
|
+
|
|
96
|
+
Sample chunks with exponentially decaying density into the past (e.g. current, $t-1, t-2, t-4, t-8\dots$):
|
|
97
|
+
|
|
98
|
+
```python
|
|
99
|
+
state_machine = DeclarativeStateMachine(chunk_spans)
|
|
100
|
+
|
|
101
|
+
@state_machine.on('log_sparse')
|
|
102
|
+
def handle_log_sparse(machine, base = 2):
|
|
103
|
+
base = int(base)
|
|
104
|
+
total = len(machine.chunk_spans)
|
|
105
|
+
|
|
106
|
+
# Powers of base distance from current chunk: 0, 1, 2, 4, 8...
|
|
107
|
+
offsets = [0] + [base ** i for i in range(10)]
|
|
108
|
+
machine.active_chunks = {
|
|
109
|
+
total - d for d in offsets if (total - d) in machine.chunk_spans
|
|
110
|
+
}
|
|
111
|
+
|
|
112
|
+
# When model emits <log_sparse> or <log_sparse base="3">
|
|
113
|
+
state_machine.step('<log_sparse base="2">')
|
|
114
|
+
|
|
115
|
+
# Reverts back to global when closed
|
|
116
|
+
state_machine.step('</log_sparse>')
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
#### 2. Sliding Window (Last $K$ Chunks)
|
|
120
|
+
|
|
121
|
+
Select the most recent $K$ context chunks via a custom `State`:
|
|
122
|
+
|
|
123
|
+
```python
|
|
124
|
+
from statemachine import State
|
|
125
|
+
from declarative_attention import DeclarativeStateMachine
|
|
126
|
+
|
|
127
|
+
class CustomAttentionMachine(DeclarativeStateMachine):
|
|
128
|
+
window_mode = State()
|
|
129
|
+
|
|
130
|
+
window = window_mode.from_(DeclarativeStateMachine.global_mode)
|
|
131
|
+
revert = DeclarativeStateMachine.global_mode.from_(
|
|
132
|
+
DeclarativeStateMachine.focus_mode,
|
|
133
|
+
DeclarativeStateMachine.local_mode,
|
|
134
|
+
window_mode
|
|
135
|
+
)
|
|
136
|
+
|
|
137
|
+
@window.on
|
|
138
|
+
def on_window(self, size = 2):
|
|
139
|
+
total = len(self.chunk_spans)
|
|
140
|
+
self.active_chunks = {total - i for i in range(int(size)) if total - i > 0}
|
|
141
|
+
|
|
142
|
+
state_machine = CustomAttentionMachine(chunk_spans)
|
|
143
|
+
|
|
144
|
+
# Model calls <window size="2">
|
|
145
|
+
state_machine.step('<window size="2">')
|
|
146
|
+
assert state_machine.is_window_mode
|
|
147
|
+
```
|
|
148
|
+
|
|
149
|
+
#### 3. Dynamic Temperature & In-Context Verification
|
|
150
|
+
|
|
151
|
+
Models can dynamically adjust sampling temperature during decode—brainstorming at high temperature, then re-attending to those thoughts at zero temperature for critical verification:
|
|
152
|
+
|
|
153
|
+
```python
|
|
154
|
+
state_machine = DeclarativeStateMachine(chunk_spans)
|
|
155
|
+
state_machine.temperature = 0.7 # default
|
|
156
|
+
|
|
157
|
+
@state_machine.on('explore')
|
|
158
|
+
def on_explore(machine, temp = 1.2):
|
|
159
|
+
machine.temperature = float(temp)
|
|
160
|
+
|
|
161
|
+
@state_machine.on('verify')
|
|
162
|
+
def on_verify(machine, temp = 0.0):
|
|
163
|
+
machine.temperature = float(temp)
|
|
164
|
+
|
|
165
|
+
# High-temperature exploration
|
|
166
|
+
state_machine.step('<explore temp="1.2">')
|
|
167
|
+
|
|
168
|
+
# Switch to zero-temperature verification on Chunk 1
|
|
169
|
+
state_machine.step('</explore><focus chunks="1"><verify temp="0.0">')
|
|
170
|
+
```
|
|
171
|
+
|
|
172
|
+
## Citations
|
|
173
|
+
|
|
174
|
+
```bibtex
|
|
175
|
+
@misc{ho2026languagemodelscontrolattention,
|
|
176
|
+
title = {Language Models Can Control Their Own Attention},
|
|
177
|
+
author = {Namgyu Ho and Huzama Ahmad and Woosung Koh and Se-Young Yun and Tal Schuster and Cicero Nogueira dos Santos},
|
|
178
|
+
year = {2026},
|
|
179
|
+
eprint = {2609.02737},
|
|
180
|
+
archivePrefix = {arXiv},
|
|
181
|
+
primaryClass = {cs.CL},
|
|
182
|
+
url = {https://arxiv.org/abs/2609.02737},
|
|
183
|
+
}
|
|
184
|
+
```
|
|
@@ -0,0 +1,139 @@
|
|
|
1
|
+
<img src="./fig1.png" width="400px"></img>
|
|
2
|
+
|
|
3
|
+
## Declarative Attention - (wip)
|
|
4
|
+
|
|
5
|
+
Implementation of the procedure in [Language Models Can Control Their Own Attention](https://arxiv.org/abs/2609.02737), from Namgyu Ho et al. of KAIST AI
|
|
6
|
+
|
|
7
|
+
## Install
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
$ pip install declarative-attention
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Usage
|
|
14
|
+
|
|
15
|
+
```python
|
|
16
|
+
import torch
|
|
17
|
+
from declarative_attention import DeclarativeStateMachine
|
|
18
|
+
|
|
19
|
+
# Define context chunk spans (1-indexed)
|
|
20
|
+
chunk_spans = [
|
|
21
|
+
(16, 1024), # Chunk 1
|
|
22
|
+
(1024, 2048), # Chunk 2
|
|
23
|
+
(2048, 3072), # Chunk 3
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
state_machine = DeclarativeStateMachine(
|
|
27
|
+
chunk_spans = chunk_spans,
|
|
28
|
+
tokenizer_decode = tokenizer.decode # or lambda t: tokenizer.decode([t])
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
# Decode loop: feed generated tokens (int, torch.long scalar, or str)
|
|
32
|
+
state_machine.step('<focus chunks="1,2">')
|
|
33
|
+
|
|
34
|
+
assert state_machine.is_focus
|
|
35
|
+
assert state_machine.active_chunks == {1, 2}
|
|
36
|
+
|
|
37
|
+
# Obtain 1D boolean attention mask for current decode step
|
|
38
|
+
mask = state_machine.get_mask(total_len = 4096) # True for kept keys, False for masked out
|
|
39
|
+
|
|
40
|
+
# Model finishes chunk reasoning and reverts back to global
|
|
41
|
+
state_machine.step('</focus>')
|
|
42
|
+
assert state_machine.is_global
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
### Custom Attention Patterns
|
|
46
|
+
|
|
47
|
+
Researchers can easily declare custom attention patterns and tags.
|
|
48
|
+
|
|
49
|
+
#### 1. Logarithmically Spaced Attention (Recent to Past)
|
|
50
|
+
|
|
51
|
+
Sample chunks with exponentially decaying density into the past (e.g. current, $t-1, t-2, t-4, t-8\dots$):
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
state_machine = DeclarativeStateMachine(chunk_spans)
|
|
55
|
+
|
|
56
|
+
@state_machine.on('log_sparse')
|
|
57
|
+
def handle_log_sparse(machine, base = 2):
|
|
58
|
+
base = int(base)
|
|
59
|
+
total = len(machine.chunk_spans)
|
|
60
|
+
|
|
61
|
+
# Powers of base distance from current chunk: 0, 1, 2, 4, 8...
|
|
62
|
+
offsets = [0] + [base ** i for i in range(10)]
|
|
63
|
+
machine.active_chunks = {
|
|
64
|
+
total - d for d in offsets if (total - d) in machine.chunk_spans
|
|
65
|
+
}
|
|
66
|
+
|
|
67
|
+
# When model emits <log_sparse> or <log_sparse base="3">
|
|
68
|
+
state_machine.step('<log_sparse base="2">')
|
|
69
|
+
|
|
70
|
+
# Reverts back to global when closed
|
|
71
|
+
state_machine.step('</log_sparse>')
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
#### 2. Sliding Window (Last $K$ Chunks)
|
|
75
|
+
|
|
76
|
+
Select the most recent $K$ context chunks via a custom `State`:
|
|
77
|
+
|
|
78
|
+
```python
|
|
79
|
+
from statemachine import State
|
|
80
|
+
from declarative_attention import DeclarativeStateMachine
|
|
81
|
+
|
|
82
|
+
class CustomAttentionMachine(DeclarativeStateMachine):
|
|
83
|
+
window_mode = State()
|
|
84
|
+
|
|
85
|
+
window = window_mode.from_(DeclarativeStateMachine.global_mode)
|
|
86
|
+
revert = DeclarativeStateMachine.global_mode.from_(
|
|
87
|
+
DeclarativeStateMachine.focus_mode,
|
|
88
|
+
DeclarativeStateMachine.local_mode,
|
|
89
|
+
window_mode
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
@window.on
|
|
93
|
+
def on_window(self, size = 2):
|
|
94
|
+
total = len(self.chunk_spans)
|
|
95
|
+
self.active_chunks = {total - i for i in range(int(size)) if total - i > 0}
|
|
96
|
+
|
|
97
|
+
state_machine = CustomAttentionMachine(chunk_spans)
|
|
98
|
+
|
|
99
|
+
# Model calls <window size="2">
|
|
100
|
+
state_machine.step('<window size="2">')
|
|
101
|
+
assert state_machine.is_window_mode
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
#### 3. Dynamic Temperature & In-Context Verification
|
|
105
|
+
|
|
106
|
+
Models can dynamically adjust sampling temperature during decode—brainstorming at high temperature, then re-attending to those thoughts at zero temperature for critical verification:
|
|
107
|
+
|
|
108
|
+
```python
|
|
109
|
+
state_machine = DeclarativeStateMachine(chunk_spans)
|
|
110
|
+
state_machine.temperature = 0.7 # default
|
|
111
|
+
|
|
112
|
+
@state_machine.on('explore')
|
|
113
|
+
def on_explore(machine, temp = 1.2):
|
|
114
|
+
machine.temperature = float(temp)
|
|
115
|
+
|
|
116
|
+
@state_machine.on('verify')
|
|
117
|
+
def on_verify(machine, temp = 0.0):
|
|
118
|
+
machine.temperature = float(temp)
|
|
119
|
+
|
|
120
|
+
# High-temperature exploration
|
|
121
|
+
state_machine.step('<explore temp="1.2">')
|
|
122
|
+
|
|
123
|
+
# Switch to zero-temperature verification on Chunk 1
|
|
124
|
+
state_machine.step('</explore><focus chunks="1"><verify temp="0.0">')
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
## Citations
|
|
128
|
+
|
|
129
|
+
```bibtex
|
|
130
|
+
@misc{ho2026languagemodelscontrolattention,
|
|
131
|
+
title = {Language Models Can Control Their Own Attention},
|
|
132
|
+
author = {Namgyu Ho and Huzama Ahmad and Woosung Koh and Se-Young Yun and Tal Schuster and Cicero Nogueira dos Santos},
|
|
133
|
+
year = {2026},
|
|
134
|
+
eprint = {2609.02737},
|
|
135
|
+
archivePrefix = {arXiv},
|
|
136
|
+
primaryClass = {cs.CL},
|
|
137
|
+
url = {https://arxiv.org/abs/2609.02737},
|
|
138
|
+
}
|
|
139
|
+
```
|
|
@@ -0,0 +1,198 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
from html.parser import HTMLParser
|
|
3
|
+
from typing import Callable, Sequence
|
|
4
|
+
|
|
5
|
+
import torch
|
|
6
|
+
from torch import Tensor
|
|
7
|
+
from statemachine import StateMachine, State
|
|
8
|
+
from torch_einops_utils import maybe
|
|
9
|
+
|
|
10
|
+
# helpers
|
|
11
|
+
|
|
12
|
+
def exists(v):
|
|
13
|
+
return v is not None
|
|
14
|
+
|
|
15
|
+
def default(v, d):
|
|
16
|
+
return v if exists(v) else d
|
|
17
|
+
|
|
18
|
+
def default_decode_fn(token: int) -> str:
|
|
19
|
+
return chr(token) if 0 <= token < 128 else ""
|
|
20
|
+
|
|
21
|
+
# types
|
|
22
|
+
|
|
23
|
+
ChunkId = int
|
|
24
|
+
StartPos = int
|
|
25
|
+
EndPos = int
|
|
26
|
+
Span = tuple[StartPos, EndPos]
|
|
27
|
+
|
|
28
|
+
# chunk spans can be:
|
|
29
|
+
# 1. dict mapping 1-indexed chunk id to token span:
|
|
30
|
+
# {1: (16, 32), 2: (32, 48)}
|
|
31
|
+
# 2. sequence of spans (auto-indexed from 1):
|
|
32
|
+
# [(16, 32), (32, 48)]
|
|
33
|
+
|
|
34
|
+
ChunkSpans = (
|
|
35
|
+
dict[ChunkId, Span] |
|
|
36
|
+
Sequence[Span]
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
TokenizerDecode = Callable[[int], str]
|
|
40
|
+
|
|
41
|
+
# general streaming tag parser
|
|
42
|
+
|
|
43
|
+
class StreamingTagParser(HTMLParser):
|
|
44
|
+
def __init__(self, on_tag: Callable[[bool, str, dict[str, str]], None]):
|
|
45
|
+
super().__init__()
|
|
46
|
+
self.on_tag = on_tag
|
|
47
|
+
|
|
48
|
+
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]):
|
|
49
|
+
self.on_tag(False, tag, dict(attrs))
|
|
50
|
+
|
|
51
|
+
def handle_endtag(self, tag: str):
|
|
52
|
+
self.on_tag(True, tag, {})
|
|
53
|
+
|
|
54
|
+
# declarative attention state machine (Ho et al., 2026)
|
|
55
|
+
|
|
56
|
+
class DeclarativeStateMachine(StateMachine):
|
|
57
|
+
"""
|
|
58
|
+
Declarative Attention State Machine
|
|
59
|
+
|
|
60
|
+
- <global> : attends to all context chunks (default)
|
|
61
|
+
- <focus chunks="1,2"> : attends only to specified chunk(s)
|
|
62
|
+
- <local> : attends to 0 context chunks
|
|
63
|
+
"""
|
|
64
|
+
|
|
65
|
+
global_mode = State(initial = True)
|
|
66
|
+
focus_mode = State()
|
|
67
|
+
local_mode = State()
|
|
68
|
+
|
|
69
|
+
focus = focus_mode.from_(global_mode, local_mode)
|
|
70
|
+
local = local_mode.from_(global_mode, focus_mode)
|
|
71
|
+
revert = global_mode.from_(focus_mode, local_mode)
|
|
72
|
+
|
|
73
|
+
def __init__(
|
|
74
|
+
self,
|
|
75
|
+
chunk_spans: ChunkSpans,
|
|
76
|
+
tokenizer_decode: TokenizerDecode | None = None
|
|
77
|
+
):
|
|
78
|
+
if isinstance(chunk_spans, (list, tuple)):
|
|
79
|
+
chunk_spans = {i + 1: span for i, span in enumerate(chunk_spans)}
|
|
80
|
+
|
|
81
|
+
self.chunk_spans = chunk_spans
|
|
82
|
+
self.active_chunks = set(chunk_spans.keys())
|
|
83
|
+
self.parser = StreamingTagParser(self._handle_tag)
|
|
84
|
+
self.tokenizer_decode = default(tokenizer_decode, default_decode_fn)
|
|
85
|
+
self.handlers: dict[str, Callable] = dict()
|
|
86
|
+
|
|
87
|
+
super().__init__()
|
|
88
|
+
|
|
89
|
+
# register custom tag handlers for arbitrary researcher logic
|
|
90
|
+
|
|
91
|
+
def on(self, tag: str, fn: Callable | None = None):
|
|
92
|
+
def decorator(handler: Callable):
|
|
93
|
+
self.handlers[tag.lower()] = handler
|
|
94
|
+
return handler
|
|
95
|
+
|
|
96
|
+
return maybe(decorator, default = decorator)(fn)
|
|
97
|
+
|
|
98
|
+
# dispatch streaming tags to transitions or custom handlers
|
|
99
|
+
|
|
100
|
+
def _handle_tag(self, is_closing: bool, tag: str, attrs: dict[str, str]):
|
|
101
|
+
tag = tag.lower()
|
|
102
|
+
|
|
103
|
+
if is_closing:
|
|
104
|
+
if hasattr(self, f"revert_{tag}"):
|
|
105
|
+
self.send(f"revert_{tag}")
|
|
106
|
+
elif not self.is_global:
|
|
107
|
+
try: self.revert()
|
|
108
|
+
except Exception: pass
|
|
109
|
+
return
|
|
110
|
+
|
|
111
|
+
# opening tag: try transition event first, then custom handler
|
|
112
|
+
|
|
113
|
+
if tag in [e.id for e in self.events]:
|
|
114
|
+
try:
|
|
115
|
+
self.send(tag, **attrs)
|
|
116
|
+
return
|
|
117
|
+
except Exception:
|
|
118
|
+
pass
|
|
119
|
+
|
|
120
|
+
if tag in self.handlers:
|
|
121
|
+
self.handlers[tag](self, **attrs)
|
|
122
|
+
|
|
123
|
+
# transition callbacks
|
|
124
|
+
|
|
125
|
+
@focus.on
|
|
126
|
+
def _on_focus(
|
|
127
|
+
self,
|
|
128
|
+
chunks: str | Sequence[int] | None = None,
|
|
129
|
+
chunk: str | int | None = None,
|
|
130
|
+
magic_chunks: str | None = None
|
|
131
|
+
):
|
|
132
|
+
target = default(chunks, default(chunk, magic_chunks))
|
|
133
|
+
|
|
134
|
+
if isinstance(target, str):
|
|
135
|
+
self.active_chunks = {int(c) for c in target.split(',') if c.strip().isdigit()}
|
|
136
|
+
elif isinstance(target, int):
|
|
137
|
+
self.active_chunks = {target}
|
|
138
|
+
elif exists(target):
|
|
139
|
+
self.active_chunks = set(target)
|
|
140
|
+
else:
|
|
141
|
+
self.active_chunks = set()
|
|
142
|
+
|
|
143
|
+
@local.on
|
|
144
|
+
def _on_local(self):
|
|
145
|
+
self.active_chunks = set()
|
|
146
|
+
|
|
147
|
+
@revert.on
|
|
148
|
+
def _on_revert(self):
|
|
149
|
+
self.active_chunks = set(self.chunk_spans.keys())
|
|
150
|
+
|
|
151
|
+
# properties
|
|
152
|
+
|
|
153
|
+
@property
|
|
154
|
+
def is_global(self):
|
|
155
|
+
return self.global_mode.is_active
|
|
156
|
+
|
|
157
|
+
@property
|
|
158
|
+
def is_focus(self):
|
|
159
|
+
return self.focus_mode.is_active
|
|
160
|
+
|
|
161
|
+
@property
|
|
162
|
+
def is_local(self):
|
|
163
|
+
return self.local_mode.is_active
|
|
164
|
+
|
|
165
|
+
@property
|
|
166
|
+
def mode(self):
|
|
167
|
+
return "global" if self.is_global else ("focus" if self.is_focus else "local")
|
|
168
|
+
|
|
169
|
+
@property
|
|
170
|
+
def is_streaming_tag(self) -> bool:
|
|
171
|
+
"""True if the parser is currently mid-tag (e.g. between '<' and '>')"""
|
|
172
|
+
return "<" in self.parser.rawdata
|
|
173
|
+
|
|
174
|
+
# token streaming (accepts int, torch LongTensor, or str)
|
|
175
|
+
|
|
176
|
+
def step(self, token: int | Tensor | str):
|
|
177
|
+
if isinstance(token, Tensor):
|
|
178
|
+
token = token.item()
|
|
179
|
+
|
|
180
|
+
text = self.tokenizer_decode(token) if isinstance(token, int) else str(token)
|
|
181
|
+
self.parser.feed(text)
|
|
182
|
+
|
|
183
|
+
# attention mask generation
|
|
184
|
+
|
|
185
|
+
def get_mask(self, total_len: int, device = None) -> Tensor:
|
|
186
|
+
mask = torch.ones(total_len, dtype = torch.bool, device = device)
|
|
187
|
+
|
|
188
|
+
if self.is_global:
|
|
189
|
+
return mask
|
|
190
|
+
|
|
191
|
+
for chunk_id, (start, end) in self.chunk_spans.items():
|
|
192
|
+
if chunk_id not in self.active_chunks:
|
|
193
|
+
mask[start:end] = False
|
|
194
|
+
|
|
195
|
+
return mask
|
|
196
|
+
|
|
197
|
+
def __call__(self, total_len: int, device = None) -> Tensor:
|
|
198
|
+
return self.get_mask(total_len, device = device)
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "declarative-attention"
|
|
3
|
+
version = "0.0.4"
|
|
4
|
+
description = "Declarative Attention - Pytorch"
|
|
5
|
+
authors = [
|
|
6
|
+
{ name = "Phil Wang", email = "lucidrains@gmail.com" }
|
|
7
|
+
]
|
|
8
|
+
readme = "README.md"
|
|
9
|
+
requires-python = ">= 3.10"
|
|
10
|
+
license = { file = "LICENSE" }
|
|
11
|
+
keywords = [
|
|
12
|
+
'artificial intelligence',
|
|
13
|
+
'deep learning',
|
|
14
|
+
'attention mechanisms'
|
|
15
|
+
]
|
|
16
|
+
|
|
17
|
+
classifiers=[
|
|
18
|
+
'Development Status :: 4 - Beta',
|
|
19
|
+
'Intended Audience :: Developers',
|
|
20
|
+
'Topic :: Scientific/Engineering :: Artificial Intelligence',
|
|
21
|
+
'License :: OSI Approved :: MIT License',
|
|
22
|
+
'Programming Language :: Python :: 3.10',
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
dependencies = [
|
|
26
|
+
"einops>=0.8.1",
|
|
27
|
+
"python-statemachine>=2.5",
|
|
28
|
+
"torch-einops-utils",
|
|
29
|
+
"torch>=2.5",
|
|
30
|
+
]
|
|
31
|
+
|
|
32
|
+
[project.urls]
|
|
33
|
+
Homepage = "https://pypi.org/project/declarative-attention/"
|
|
34
|
+
Repository = "https://github.com/lucidrains/declarative-attention"
|
|
35
|
+
|
|
36
|
+
[project.optional-dependencies]
|
|
37
|
+
examples = []
|
|
38
|
+
test = [
|
|
39
|
+
"pytest"
|
|
40
|
+
]
|
|
41
|
+
|
|
42
|
+
[tool.pytest.ini_options]
|
|
43
|
+
pythonpath = [
|
|
44
|
+
"."
|
|
45
|
+
]
|
|
46
|
+
|
|
47
|
+
[build-system]
|
|
48
|
+
requires = ["hatchling"]
|
|
49
|
+
build-backend = "hatchling.build"
|
|
50
|
+
|
|
51
|
+
[tool.rye]
|
|
52
|
+
managed = true
|
|
53
|
+
dev-dependencies = []
|
|
54
|
+
|
|
55
|
+
[tool.hatch.metadata]
|
|
56
|
+
allow-direct-references = true
|
|
57
|
+
|
|
58
|
+
[tool.hatch.build]
|
|
59
|
+
include = [
|
|
60
|
+
"declarative_attention",
|
|
61
|
+
"pyproject.toml",
|
|
62
|
+
"README.md",
|
|
63
|
+
"LICENSE"
|
|
64
|
+
]
|
|
65
|
+
|
|
66
|
+
[tool.hatch.build.targets.wheel]
|
|
67
|
+
packages = ["declarative_attention"]
|