ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
Browse files- .gitignore +45 -0
- LICENSE +202 -0
- README.md +377 -0
- docs/AUDIT.md +62 -0
- docs/ZeroPoint.md +247 -0
- docs/ZeroPoint_EN.md +214 -0
- experiments/ablation.md +21 -0
- experiments/ablation.py +86 -0
- experiments/common.py +119 -0
- experiments/compare.md +53 -0
- experiments/compare.py +128 -0
- experiments/diagnose.md +65 -0
- experiments/diagnose.py +216 -0
- experiments/ema_longrun.md +55 -0
- experiments/ema_longrun.py +128 -0
- experiments/recall.md +37 -0
- experiments/recall.py +87 -0
- experiments/recall_diag.py +118 -0
- experiments/recall_prereg.md +13 -0
- experiments/results/20261011-1223-compare.md +42 -0
- experiments/results/20261011-1258-ablation.md +22 -0
- experiments/results/20261011-1305-recall.md +12 -0
- experiments/results/20261018-0838-ema_longrun.md +57 -0
- experiments/results/20261018-0901-diagnose.md +67 -0
- experiments/results/20261101-0814-slots.md +37 -0
- experiments/results/20261101-0936-recall.md +12 -0
- experiments/slots.md +35 -0
- experiments/slots.py +88 -0
- models/__init__.py +4 -0
- models/transformer.py +63 -0
- models/zeropoint.py +217 -0
- requirements.txt +1 -0
- scripts/upload.md +92 -0
- tests/golden/zeropoint_ref.npz +3 -0
- tests/make_golden.py +32 -0
- tests/test_zerocg.py +250 -0
- tests/test_zeropoint.py +153 -0
- zerocg/__init__.py +11 -0
- zerocg/functional.py +69 -0
- zerocg/nn.py +162 -0
- zerocg/optim.py +87 -0
- zerocg/tensor.py +320 -0
.gitignore
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.py[cod]
|
| 4 |
+
*$py.class
|
| 5 |
+
*.so
|
| 6 |
+
.Python
|
| 7 |
+
build/
|
| 8 |
+
dist/
|
| 9 |
+
*.egg-info/
|
| 10 |
+
.venv/
|
| 11 |
+
venv/
|
| 12 |
+
env/
|
| 13 |
+
.pytest_cache/
|
| 14 |
+
.mypy_cache/
|
| 15 |
+
.ruff_cache/
|
| 16 |
+
.ipynb_checkpoints/
|
| 17 |
+
|
| 18 |
+
# macOS
|
| 19 |
+
.DS_Store
|
| 20 |
+
|
| 21 |
+
# Model weights and checkpoints (this release ships no weights)
|
| 22 |
+
*.npz
|
| 23 |
+
*.npy
|
| 24 |
+
!tests/golden/*.npz
|
| 25 |
+
*.pt
|
| 26 |
+
*.pth
|
| 27 |
+
*.bin
|
| 28 |
+
*.safetensors
|
| 29 |
+
*.ckpt
|
| 30 |
+
*.pkl
|
| 31 |
+
*.pickle
|
| 32 |
+
checkpoints/
|
| 33 |
+
weights/
|
| 34 |
+
runs/
|
| 35 |
+
out/
|
| 36 |
+
|
| 37 |
+
# Data and training artifacts
|
| 38 |
+
corpus/
|
| 39 |
+
data/
|
| 40 |
+
datasets/
|
| 41 |
+
*.log
|
| 42 |
+
|
| 43 |
+
# Secrets
|
| 44 |
+
.env
|
| 45 |
+
.env.*
|
LICENSE
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Apache License
|
| 2 |
+
Version 2.0, January 2004
|
| 3 |
+
http://www.apache.org/licenses/
|
| 4 |
+
|
| 5 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
| 6 |
+
|
| 7 |
+
1. Definitions.
|
| 8 |
+
|
| 9 |
+
"License" shall mean the terms and conditions for use, reproduction,
|
| 10 |
+
and distribution as defined by Sections 1 through 9 of this document.
|
| 11 |
+
|
| 12 |
+
"Licensor" shall mean the copyright owner or entity authorized by
|
| 13 |
+
the copyright owner that is granting the License.
|
| 14 |
+
|
| 15 |
+
"Legal Entity" shall mean the union of the acting entity and all
|
| 16 |
+
other entities that control, are controlled by, or are under common
|
| 17 |
+
control with that entity. For the purposes of this definition,
|
| 18 |
+
"control" means (i) the power, direct or indirect, to cause the
|
| 19 |
+
direction or management of such entity, whether by contract or
|
| 20 |
+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
| 21 |
+
outstanding shares, or (iii) beneficial ownership of such entity.
|
| 22 |
+
|
| 23 |
+
"You" (or "Your") shall mean an individual or Legal Entity
|
| 24 |
+
exercising permissions granted by this License.
|
| 25 |
+
|
| 26 |
+
"Source" form shall mean the preferred form for making modifications,
|
| 27 |
+
including but not limited to software source code, documentation
|
| 28 |
+
source, and configuration files.
|
| 29 |
+
|
| 30 |
+
"Object" form shall mean any form resulting from mechanical
|
| 31 |
+
transformation or translation of a Source form, including but
|
| 32 |
+
not limited to compiled object code, generated documentation,
|
| 33 |
+
and conversions to other media types.
|
| 34 |
+
|
| 35 |
+
"Work" shall mean the work of authorship, whether in Source or
|
| 36 |
+
Object form, made available under the License, as indicated by a
|
| 37 |
+
copyright notice that is included in or attached to the work
|
| 38 |
+
(an example is provided in the Appendix below).
|
| 39 |
+
|
| 40 |
+
"Derivative Works" shall mean any work, whether in Source or Object
|
| 41 |
+
form, that is based on (or derived from) the Work and for which the
|
| 42 |
+
editorial revisions, annotations, elaborations, or other modifications
|
| 43 |
+
represent, as a whole, an original work of authorship. For the purposes
|
| 44 |
+
of this License, Derivative Works shall not include works that remain
|
| 45 |
+
separable from, or merely link (or bind by name) to the interfaces of,
|
| 46 |
+
the Work and Derivative Works thereof.
|
| 47 |
+
|
| 48 |
+
"Contribution" shall mean any work of authorship, including
|
| 49 |
+
the original version of the Work and any modifications or additions
|
| 50 |
+
to that Work or Derivative Works thereof, that is intentionally
|
| 51 |
+
submitted to Licensor for inclusion in the Work by the copyright owner
|
| 52 |
+
or by an individual or Legal Entity authorized to submit on behalf of
|
| 53 |
+
the copyright owner. For the purposes of this definition, "submitted"
|
| 54 |
+
means any form of electronic, verbal, or written communication sent
|
| 55 |
+
to the Licensor or its representatives, including but not limited to
|
| 56 |
+
communication on electronic mailing lists, source code control systems,
|
| 57 |
+
and issue tracking systems that are managed by, or on behalf of, the
|
| 58 |
+
Licensor for the purpose of discussing and improving the Work, but
|
| 59 |
+
excluding communication that is conspicuously marked or otherwise
|
| 60 |
+
designated in writing by the copyright owner as "Not a Contribution."
|
| 61 |
+
|
| 62 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity
|
| 63 |
+
on behalf of whom a Contribution has been received by Licensor and
|
| 64 |
+
subsequently incorporated within the Work.
|
| 65 |
+
|
| 66 |
+
2. Grant of Copyright License. Subject to the terms and conditions of
|
| 67 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 68 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 69 |
+
copyright license to reproduce, prepare Derivative Works of,
|
| 70 |
+
publicly display, publicly perform, sublicense, and distribute the
|
| 71 |
+
Work and such Derivative Works in Source or Object form.
|
| 72 |
+
|
| 73 |
+
3. Grant of Patent License. Subject to the terms and conditions of
|
| 74 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 75 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 76 |
+
(except as stated in this section) patent license to make, have made,
|
| 77 |
+
use, offer to sell, sell, import, and otherwise transfer the Work,
|
| 78 |
+
where such license applies only to those patent claims licensable
|
| 79 |
+
by such Contributor that are necessarily infringed by their
|
| 80 |
+
Contribution(s) alone or by combination of their Contribution(s)
|
| 81 |
+
with the Work to which such Contribution(s) was submitted. If You
|
| 82 |
+
institute patent litigation against any entity (including a
|
| 83 |
+
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
| 84 |
+
or a Contribution incorporated within the Work constitutes direct
|
| 85 |
+
or contributory patent infringement, then any patent licenses
|
| 86 |
+
granted to You under this License for that Work shall terminate
|
| 87 |
+
as of the date such litigation is filed.
|
| 88 |
+
|
| 89 |
+
4. Redistribution. You may reproduce and distribute copies of the
|
| 90 |
+
Work or Derivative Works thereof in any medium, with or without
|
| 91 |
+
modifications, and in Source or Object form, provided that You
|
| 92 |
+
meet the following conditions:
|
| 93 |
+
|
| 94 |
+
(a) You must give any other recipients of the Work or
|
| 95 |
+
Derivative Works a copy of this License; and
|
| 96 |
+
|
| 97 |
+
(b) You must cause any modified files to carry prominent notices
|
| 98 |
+
stating that You changed the files; and
|
| 99 |
+
|
| 100 |
+
(c) You must retain, in the Source form of any Derivative Works
|
| 101 |
+
that You distribute, all copyright, patent, trademark, and
|
| 102 |
+
attribution notices from the Source form of the Work,
|
| 103 |
+
excluding those notices that do not pertain to any part of
|
| 104 |
+
the Derivative Works; and
|
| 105 |
+
|
| 106 |
+
(d) If the Work includes a "NOTICE" text file as part of its
|
| 107 |
+
distribution, then any Derivative Works that You distribute must
|
| 108 |
+
include a readable copy of the attribution notices contained
|
| 109 |
+
within such NOTICE file, excluding those notices that do not
|
| 110 |
+
pertain to any part of the Derivative Works, in at least one
|
| 111 |
+
of the following places: within a NOTICE text file distributed
|
| 112 |
+
as part of the Derivative Works; within the Source form or
|
| 113 |
+
documentation, if provided along with the Derivative Works; or,
|
| 114 |
+
within a display generated by the Derivative Works, if and
|
| 115 |
+
wherever such third-party notices normally appear. The contents
|
| 116 |
+
of the NOTICE file are for informational purposes only and
|
| 117 |
+
do not modify the License. You may add Your own attribution
|
| 118 |
+
notices within Derivative Works that You distribute, alongside
|
| 119 |
+
or as an addendum to the NOTICE text from the Work, provided
|
| 120 |
+
that such additional attribution notices cannot be construed
|
| 121 |
+
as modifying the License.
|
| 122 |
+
|
| 123 |
+
You may add Your own copyright statement to Your modifications and
|
| 124 |
+
may provide additional or different license terms and conditions
|
| 125 |
+
for use, reproduction, or distribution of Your modifications, or
|
| 126 |
+
for any such Derivative Works as a whole, provided Your use,
|
| 127 |
+
reproduction, and distribution of the Work otherwise complies with
|
| 128 |
+
the conditions stated in this License.
|
| 129 |
+
|
| 130 |
+
5. Submission of Contributions. Unless You explicitly state otherwise,
|
| 131 |
+
any Contribution intentionally submitted for inclusion in the Work
|
| 132 |
+
by You to the Licensor shall be under the terms and conditions of
|
| 133 |
+
this License, without any additional terms or conditions.
|
| 134 |
+
Notwithstanding the above, nothing herein shall supersede or modify
|
| 135 |
+
the terms of any separate license agreement you may have executed
|
| 136 |
+
with Licensor regarding such Contributions.
|
| 137 |
+
|
| 138 |
+
6. Trademarks. This License does not grant permission to use the trade
|
| 139 |
+
names, trademarks, service marks, or product names of the Licensor,
|
| 140 |
+
except as required for reasonable and customary use in describing the
|
| 141 |
+
origin of the Work and reproducing the content of the NOTICE file.
|
| 142 |
+
|
| 143 |
+
7. Disclaimer of Warranty. Unless required by applicable law or
|
| 144 |
+
agreed to in writing, Licensor provides the Work (and each
|
| 145 |
+
Contributor provides its Contributions) on an "AS IS" BASIS,
|
| 146 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
| 147 |
+
implied, including, without limitation, any warranties or conditions
|
| 148 |
+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
| 149 |
+
PARTICULAR PURPOSE. You are solely responsible for determining the
|
| 150 |
+
appropriateness of using or redistributing the Work and assume any
|
| 151 |
+
risks associated with Your exercise of permissions under this License.
|
| 152 |
+
|
| 153 |
+
8. Limitation of Liability. In no event and under no legal theory,
|
| 154 |
+
whether in tort (including negligence), contract, or otherwise,
|
| 155 |
+
unless required by applicable law (such as deliberate and grossly
|
| 156 |
+
negligent acts) or agreed to in writing, shall any Contributor be
|
| 157 |
+
liable to You for damages, including any direct, indirect, special,
|
| 158 |
+
incidental, or consequential damages of any character arising as a
|
| 159 |
+
result of this License or out of the use or inability to use the
|
| 160 |
+
Work (including but not limited to damages for loss of goodwill,
|
| 161 |
+
work stoppage, computer failure or malfunction, or any and all
|
| 162 |
+
other commercial damages or losses), even if such Contributor
|
| 163 |
+
has been advised of the possibility of such damages.
|
| 164 |
+
|
| 165 |
+
9. Accepting Warranty or Additional Liability. While redistributing
|
| 166 |
+
the Work or Derivative Works thereof, You may choose to offer,
|
| 167 |
+
and charge a fee for, acceptance of support, warranty, indemnity,
|
| 168 |
+
or other liability obligations and/or rights consistent with this
|
| 169 |
+
License. However, in accepting such obligations, You may act only
|
| 170 |
+
on Your own behalf and on Your sole responsibility, not on behalf
|
| 171 |
+
of any other Contributor, and only if You agree to indemnify,
|
| 172 |
+
defend, and hold each Contributor harmless for any liability
|
| 173 |
+
incurred by, or claims asserted against, such Contributor by reason
|
| 174 |
+
of your accepting any such warranty or additional liability.
|
| 175 |
+
|
| 176 |
+
END OF TERMS AND CONDITIONS
|
| 177 |
+
|
| 178 |
+
APPENDIX: How to apply the Apache License to your work.
|
| 179 |
+
|
| 180 |
+
To apply the Apache License to your work, attach the following
|
| 181 |
+
boilerplate notice, with the fields enclosed by brackets "{}"
|
| 182 |
+
replaced with your own identifying information. (Don't include
|
| 183 |
+
the brackets!) The text should be enclosed in the appropriate
|
| 184 |
+
comment syntax for the file format. We also recommend that a
|
| 185 |
+
file or class name and description of purpose be included on the
|
| 186 |
+
same "printed page" as the copyright notice for easier
|
| 187 |
+
identification within third-party archives.
|
| 188 |
+
|
| 189 |
+
Copyright 2026 ZeroPoint contributors
|
| 190 |
+
|
| 191 |
+
Licensed under the Apache License, Version 2.0 (the "License");
|
| 192 |
+
you may not use this file except in compliance with the License.
|
| 193 |
+
You may obtain a copy of the License at
|
| 194 |
+
|
| 195 |
+
http://www.apache.org/licenses/LICENSE-2.0
|
| 196 |
+
|
| 197 |
+
Unless required by applicable law or agreed to in writing, software
|
| 198 |
+
distributed under the License is distributed on an "AS IS" BASIS,
|
| 199 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 200 |
+
See the License for the specific language governing permissions and
|
| 201 |
+
limitations under the License.
|
| 202 |
+
|
README.md
ADDED
|
@@ -0,0 +1,377 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
tags:
|
| 6 |
+
- zeropoint
|
| 7 |
+
- zerocg
|
| 8 |
+
- numpy
|
| 9 |
+
- from-scratch
|
| 10 |
+
- sequence-modeling
|
| 11 |
+
- memory-slots
|
| 12 |
+
- linear-attention
|
| 13 |
+
- research
|
| 14 |
+
pretty_name: ZeroPoint
|
| 15 |
+
---
|
| 16 |
+
|
| 17 |
+
# ZeroPoint — a new, from-scratch framework for sequence modeling
|
| 18 |
+
|
| 19 |
+
**Status: v0 research prototype (November 2026). This is a brand-new framework: it is not a
|
| 20 |
+
fine-tune, not a wrapper around PyTorch / `transformers`, and not a derivative of an existing
|
| 21 |
+
released model. Both the framework and the architecture are implemented here from scratch,
|
| 22 |
+
in pure Python + NumPy.**
|
| 23 |
+
|
| 24 |
+
ZeroPoint is two things:
|
| 25 |
+
|
| 26 |
+
1. **ZeroCG** — a small deep-learning framework written from scratch: `Tensor` with
|
| 27 |
+
reverse-mode automatic differentiation, `nn` layers, AdamW (checkpointable), and
|
| 28 |
+
matrix-multiply FLOP counting. Dependencies: Python 3 + NumPy. Nothing else.
|
| 29 |
+
2. **ZeroPoint** — a sequence architecture that replaces self-attention with a set of
|
| 30 |
+
**memory slots that forget at different, learnable speeds**. Training uses a
|
| 31 |
+
**chunked-parallel** form (chunk = 32); inference uses a **per-character recurrence**
|
| 32 |
+
whose state size is fixed and independent of text length, and it needs **no positional
|
| 33 |
+
encoding**.
|
| 34 |
+
|
| 35 |
+
**This repository contains no pretrained weights.** It is an architecture + framework
|
| 36 |
+
release, together with the full experiment suite that compares ZeroPoint against a
|
| 37 |
+
GPT-style Transformer baseline at ~115k parameters and a character level. A separate
|
| 38 |
+
1.4M-parameter model trained on a private corpus (internal name *ZeroX*) is **not** part
|
| 39 |
+
of this release, and neither are its weights, corpus, or training runs.
|
| 40 |
+
|
| 41 |
+
---
|
| 42 |
+
|
| 43 |
+
## TL;DR — the measured comparison
|
| 44 |
+
|
| 45 |
+
On the only head-to-head experiment (same data, same parameter budget, same training
|
| 46 |
+
recipe, same evaluation), **ZeroPoint and the Transformer are statistically tied** under
|
| 47 |
+
the project's own "difference > 2× seed std" rule. We do **not** claim ZeroPoint is better.
|
| 48 |
+
|
| 49 |
+
| model | params | val loss (nats/char, 3 seeds) | matmul FLOPs/char @T=64 | inference state @64 / @1024 / @4096 (floats) |
|
| 50 |
+
|---|---|---|---|---|
|
| 51 |
+
| Transformer (learned PE) | 116,576 | 1.5452 ± 0.0100 | 241,664 | 16,384 / 262,144 / 1,048,576 |
|
| 52 |
+
| Transformer (sinusoidal PE) | 112,480 | 1.5271 ± 0.0083 | 241,664 | 16,384 / 262,144 / 1,048,576 |
|
| 53 |
+
| **ZeroPoint** | **115,240** | **1.5214 ± 0.0134** | **232,448** | **2,048 / 2,048 / 2,048** |
|
| 54 |
+
|
| 55 |
+
- ZeroPoint − Transformer = 0.0238; 2 × max seed std = 0.0268 → **ratio 0.89, not
|
| 56 |
+
significant**. ZeroPoint − Transformer(sin) = 0.0057 → ratio 0.21, not significant.
|
| 57 |
+
- The transformer's KV cache is **8×** ZeroPoint's state at 64 characters, **128×** at
|
| 58 |
+
1,024 characters and **512×** at 4,096 characters. ZeroPoint's state is 2,048 floats
|
| 59 |
+
(8 KiB in fp32) at any length.
|
| 60 |
+
- The Transformer's per-character matmul cost grows with context (2,306,048 FLOPs/char at
|
| 61 |
+
4,096) while ZeroPoint's stays ~constant (234,464). **Wall-clock speed was not measured**
|
| 62 |
+
— this release ships FLOP counts, not throughput numbers.
|
| 63 |
+
|
| 64 |
+
Read [**Comparison conditions**](#comparison-conditions-read-before-quoting-any-number)
|
| 65 |
+
before quoting any number, and [**Limitations**](#limitations) before drawing conclusions.
|
| 66 |
+
|
| 67 |
+
---
|
| 68 |
+
|
| 69 |
+
## What is actually new (and what is not)
|
| 70 |
+
|
| 71 |
+
**New code.** `zerocg/` is a self-contained NumPy implementation: reverse-mode autograd,
|
| 72 |
+
layers, AdamW, checkpointing, FLOP counting. `models/zeropoint.py` is the architecture
|
| 73 |
+
implementation. Neither depends on PyTorch; the early PyTorch reference scripts of the
|
| 74 |
+
project are not part of this release and were never used for the results.
|
| 75 |
+
|
| 76 |
+
**New combination.** The specific design — **per-slot learnable decay λ with a geometric
|
| 77 |
+
half-life initialisation (1.5 → 64 characters) + softmax competition when writing +
|
| 78 |
+
a write-strength gate + chunked-parallel / recurrent dual form + no positional encoding +
|
| 79 |
+
learnable initial slot state** — is, to our knowledge, a new combination. It is the object
|
| 80 |
+
of study here.
|
| 81 |
+
|
| 82 |
+
**Not new.** Every individual ingredient has prior art, and the design document says so:
|
| 83 |
+
memory networks / Neural Turing Machines (soft read-write memory), linear attention
|
| 84 |
+
(Katharopoulos et al., 2020 — parallel/recurrent dual form), RetNet (Sun et al., 2023 —
|
| 85 |
+
decayed recurrence without positional encoding), and Mamba / state-space models (fixed-size
|
| 86 |
+
recurrent state). This work makes **no priority claim** and is not "the first" anything.
|
| 87 |
+
It is also not a claim that memory slots are generally better than attention: at this scale
|
| 88 |
+
they tie.
|
| 89 |
+
|
| 90 |
+
---
|
| 91 |
+
|
| 92 |
+
## Comparison conditions (read before quoting any number)
|
| 93 |
+
|
| 94 |
+
Everything below is what the two sides of the comparison actually shared, measured on
|
| 95 |
+
**2026-10-11** (head-to-head and ablation) and **2026-10-18** (length re-evaluation).
|
| 96 |
+
|
| 97 |
+
| item | setting |
|
| 98 |
+
|---|---|
|
| 99 |
+
| Corpus | ASCII-only Python standard-library `.py` source, **1,558,395** characters, vocabulary **96** |
|
| 100 |
+
| Split | contiguous 90 % train / 10 % validation |
|
| 101 |
+
| Task | character-level next-character prediction |
|
| 102 |
+
| Models | width 64, 2 layers, 4 heads (D = 16), MLP hidden 256 (4×), Pre-LN, GELU, no dropout, no weight tying |
|
| 103 |
+
| ZeroPoint specifics | 16 slots/head, initial half-lives geometric 1.5→64 chars, learnable λ and s0, chunk 32, no positional encoding |
|
| 104 |
+
| Context length | 64 characters (training and short-context eval) |
|
| 105 |
+
| Batch | 32 sequences × 64 characters = 2,048 chars/step |
|
| 106 |
+
| Steps | 2,000 (about 2 minutes per model on one CPU core) |
|
| 107 |
+
| Optimizer | AdamW β = (0.9, 0.999), weight decay 0.01, gradient clip 1.0 |
|
| 108 |
+
| LR schedule | 100-step warmup, cosine decay to 10 % |
|
| 109 |
+
| LR selection | per model: sweep {1e-3, 3e-3, 6e-3} on seed 0, take the best for seeds 1–2; all three models selected 6e-3 |
|
| 110 |
+
| Seeds | 0, 1, 2 (seed controls initialisation and data order) |
|
| 111 |
+
| Evaluation | mean cross-entropy over 40 fixed validation batches × 32 sequences × 64 chars (seed 999); mean ± sample std over 3 seeds |
|
| 112 |
+
| Hardware | one Apple Silicon Mac, CPU only, NumPy float32 — no GPU |
|
| 113 |
+
| Data used | only the reports in `experiments/`; the numbers below were independently re-derived from saved checkpoints where possible |
|
| 114 |
+
|
| 115 |
+
**Not controlled / things this comparison does not tell you:**
|
| 116 |
+
|
| 117 |
+
- **One corpus, one domain, one split.** Source code is not natural language. No downstream
|
| 118 |
+
tasks, no instruction following, no reasoning evals — only validation loss.
|
| 119 |
+
- **Tiny scale.** 115k parameters, 2 layers, 64-char context, 2,000 steps. Nothing here
|
| 120 |
+
extrapolates to large models or long training.
|
| 121 |
+
- **The baseline is a minimal GPT-style decoder** with learned or sinusoidal positional
|
| 122 |
+
encoding, Pre-LN and GELU. It is not tuned beyond a 3-value LR sweep; no RoPE/ALiBi, no
|
| 123 |
+
dropout, no weight decay tuning.
|
| 124 |
+
- **3 seeds only**, and the LR is selected on seed 0, which is also included in the 3-seed
|
| 125 |
+
mean (a winner's-curse of ≤ ~0.01 nats/char; the same protocol was used for all models).
|
| 126 |
+
- **Transformer checkpoints were not saved**, so the Transformer rows cannot be re-derived
|
| 127 |
+
from artifacts — they rest on the original report. The ZeroPoint side *was* re-derived
|
| 128 |
+
bit-for-bit (see [Data verification](#data-verification)).
|
| 129 |
+
- **No wall-clock or memory benchmark.** FLOPs count matrix multiplications only; CPU
|
| 130 |
+
wall-clock is dominated by Python overhead and would not be a meaningful comparison.
|
| 131 |
+
- **Length extrapolation is evaluated at two different sample sizes** (32 sequences for
|
| 132 |
+
ZeroPoint, 8 for the sinusoidal Transformer — the baseline checkpoints are gone).
|
| 133 |
+
- **Ablation numbers use a different fixed LR** (3e-3) than the head-to-head (6e-3), so the
|
| 134 |
+
two tables are only internally comparable. Do not rank across them.
|
| 135 |
+
- **Exact recall is not measured** by this release: the recall experiment is inconclusive
|
| 136 |
+
because of its learning-rate configuration (see below), and must not be cited as an
|
| 137 |
+
architecture comparison.
|
| 138 |
+
|
| 139 |
+
---
|
| 140 |
+
|
| 141 |
+
## Head-to-head results
|
| 142 |
+
|
| 143 |
+
Protocol: identical width/layer/head/MLP configuration; parameter counts within ~4 %
|
| 144 |
+
(115,240 / 116,576 / 112,480); identical data, batches, seeds, optimiser, schedule and
|
| 145 |
+
evaluation; the only difference is the token-mixing layer (memory slots vs self-attention)
|
| 146 |
+
and the positional encoding.
|
| 147 |
+
|
| 148 |
+
```
|
| 149 |
+
validation loss (nats/char, lower is better; 3 seeds, 2000 steps, lr 6e-3)
|
| 150 |
+
ZeroPoint 1.5214 ± 0.0134
|
| 151 |
+
Transformer (sin) 1.5271 ± 0.0083
|
| 152 |
+
Transformer (learned) 1.5452 ± 0.0100
|
| 153 |
+
```
|
| 154 |
+
|
| 155 |
+
Per-seed ZeroPoint values: 1.5311 / 1.5307 / 1.5024. The largest gap (vs learned-PE
|
| 156 |
+
Transformer) is 0.0238 nats/char, **below the project's 2 × seed-std decision line of
|
| 157 |
+
0.0268**. Under a stricter standard-error criterion the gap would be ~2.5σ in ZeroPoint's
|
| 158 |
+
favour, but that is not the rule this project pre-registered, and 3 seeds is thin. **The
|
| 159 |
+
honest reading is: indistinguishable at this scale.**
|
| 160 |
+
|
| 161 |
+
### Inference state and compute
|
| 162 |
+
|
| 163 |
+
| | @64 chars | @1,024 chars | @4,096 chars |
|
| 164 |
+
|---|---|---|---|
|
| 165 |
+
| Transformer KV cache (floats) | 16,384 | 262,144 (128×) | 1,048,576 (512×) |
|
| 166 |
+
| ZeroPoint state (floats) | 2,048 | 2,048 | 2,048 |
|
| 167 |
+
| Transformer matmul FLOPs/char | 241,664 | 733,184 | 2,306,048 |
|
| 168 |
+
| ZeroPoint matmul FLOPs/char | 232,448 | 234,368 | 234,464 |
|
| 169 |
+
|
| 170 |
+
ZeroPoint's state is per-layer slots (2 layers × 4 heads × 16 slots × 16 dims) and is
|
| 171 |
+
independent of sequence length. This is the property the architecture was designed for: a
|
| 172 |
+
fixed-size recurrent state with no KV cache. It is **not** a claim that this is faster in
|
| 173 |
+
wall-clock time — see the limitations.
|
| 174 |
+
|
| 175 |
+
---
|
| 176 |
+
|
| 177 |
+
## Length behaviour (corrected 32-sequence evaluation)
|
| 178 |
+
|
| 179 |
+
Trained on 64 characters, then evaluated on 4,096-character sequences, in segments. 3
|
| 180 |
+
seeds, 32 validation sequences, same chunked-parallel algorithm for every model:
|
| 181 |
+
|
| 182 |
+
| positions | ZeroPoint (trained at 64) | ZeroPoint (trained at 256) |
|
| 183 |
+
|---|---|---|
|
| 184 |
+
| 0–63 | 1.4372 ± 0.0187 | 1.4741 ± 0.0117 |
|
| 185 |
+
| 64–255 (4× training length) | 1.4742 ± 0.0341 | 1.4486 ± 0.0098 |
|
| 186 |
+
| 256–1023 | 1.6617 ± 0.0234 | 1.5098 ± 0.0022 |
|
| 187 |
+
| 1024–4095 (64× training length) | 1.6932 ± 0.0305 | 1.5155 ± 0.0048 |
|
| 188 |
+
| far − near gap | **+0.2190 ± 0.0630** | **+0.0669 ± 0.0051** |
|
| 189 |
+
|
| 190 |
+
- At 4× the training length the loss rises by **0.037** nats/char; at 64× by **0.219**.
|
| 191 |
+
Extrapolation degrades gracefully, but it **does** degrade — do not quote "handles 4×
|
| 192 |
+
training length" without the +0.037.
|
| 193 |
+
- Training at 256 characters reduces the far−near gap 3.3× (to +0.067) but costs +0.036 at
|
| 194 |
+
64 characters. That run also changed batch size 32 → 8 (to keep 2,048 chars/step), so the
|
| 195 |
+
short-range cost is confounded with batch size; treat it as a hypothesis, not isolated.
|
| 196 |
+
- **Sinusoidal-PE Transformer collapses outside the training window:** 1.4512 within
|
| 197 |
+
0–63, then **4.6359** at 64–255 and **4.6057** at 256–1023, versus ln 96 = 4.5642, i.e.
|
| 198 |
+
chance level. This is an 8-sequence × 3-seed measurement, and the baseline checkpoints
|
| 199 |
+
were not retained, so it cannot be re-run with 32 sequences. It is a real, large effect,
|
| 200 |
+
but it is also a known weakness of plain sinusoidal absolute PE — **it is not evidence
|
| 201 |
+
that attention in general cannot extrapolate** (relative-position and NoPE models are
|
| 202 |
+
outside this comparison).
|
| 203 |
+
|
| 204 |
+
The earlier 8-sequence values (1.5550 / 1.6164 / 1.6410 / 1.7040) are kept only in the
|
| 205 |
+
archived report and are **superseded** by the table above; the 8-sequence sample was
|
| 206 |
+
noise-dominated (the same model scores 1.5550 vs 1.4372 on 0–63 depending on sample size).
|
| 207 |
+
|
| 208 |
+
---
|
| 209 |
+
|
| 210 |
+
## Ablations (fixed LR 3e-3; internally comparable only)
|
| 211 |
+
|
| 212 |
+
Each row changes one thing; parameter counts are matched within 1 % by adjusting the MLP
|
| 213 |
+
width; 2,000 steps, 3 seeds, same corpus. Base: **1.5500 ± 0.0073**. Decision line per
|
| 214 |
+
row: 2 × max(base std, row std). Δ is row − base (positive = worse).
|
| 215 |
+
|
| 216 |
+
| variant | Δ (nats/char) | Δ / threshold | verdict |
|
| 217 |
+
|---|---|---|---|
|
| 218 |
+
| uniform half-life init (no hierarchy) | **+0.0510** | 3.23× | clear: hierarchy matters |
|
| 219 |
+
| + sinusoidal positional encoding | **+0.0606** | 3.88× | clear: no PE is better here |
|
| 220 |
+
| λ frozen at initial values | **+0.0284** | 1.46× | clear: learnable decay matters |
|
| 221 |
+
| 1 head (64 dims, MLP widened to match) | **+0.0276** | 1.89× | clear: multiple heads matter |
|
| 222 |
+
| 8 heads (8 dims each, MLP narrowed) | +0.0173 | 1.19× | marginal / weak |
|
| 223 |
+
| 4 slots/head | −0.0034 | 0.20× | no significant difference |
|
| 224 |
+
| 32 slots/head | +0.0124 | 0.85× | no significant difference |
|
| 225 |
+
| no write-strength gate | +0.0022 | 0.15× | no significant difference |
|
| 226 |
+
| sigmoid write instead of softmax competition | −0.0013 | 0.09× | no significant difference |
|
| 227 |
+
|
| 228 |
+
At this scale, the things that **do** matter are the hierarchical half-life initialisation,
|
| 229 |
+
learnable λ, multiple heads and the absence of positional encoding. Slot count (4–32),
|
| 230 |
+
the write-strength gate, and write competition do **not** show a significant effect here.
|
| 231 |
+
Note the base loss differs from the head-to-head table because the ablation fixes
|
| 232 |
+
lr = 3e-3 (base 1.5500) while the head-to-head uses the selected lr = 6e-3 (1.5214).
|
| 233 |
+
|
| 234 |
+
---
|
| 235 |
+
|
| 236 |
+
## Exact recall: inconclusive — do not cite as an architecture comparison
|
| 237 |
+
|
| 238 |
+
The long-distance associative-recall experiment (`k1 v1 … kN vN [query kq] → vq`) was meant
|
| 239 |
+
to test the expected weakness of a fixed slot state. It produced **no usable conclusion**:
|
| 240 |
+
|
| 241 |
+
- With the original constant lr = 3e-3, both architectures plateau near 50 % on N=2
|
| 242 |
+
(chance is 6 % for 16 values) and near 16–23 % for larger N, after 20,000 steps. The
|
| 243 |
+
plateau is a **training-configuration** artefact, not an architecture property: with a
|
| 244 |
+
1e-3 cosine schedule, a Transformer trained on N=2 alone reaches **100 %** in 8,000
|
| 245 |
+
steps, and the diagnostic shows the loss ≈ ln 2 (50/50 hedge) trap.
|
| 246 |
+
- The mixed-N setup also dilutes the easy examples (N=2 is ~1/15 of the data), so the
|
| 247 |
+
experiment needs a curriculum and/or per-N weighting before it can compare architectures.
|
| 248 |
+
|
| 249 |
+
The 20,000-step table (Transformer 50/30/23/20/18 %; ZeroPoint-16 51/32/23/20/16 %;
|
| 250 |
+
ZeroPoint-32 49/31/24/20/17 %) is therefore **not** evidence that ZeroPoint is worse at
|
| 251 |
+
recall. It is a negative result about the experimental setup. A corrected re-run is future
|
| 252 |
+
work.
|
| 253 |
+
|
| 254 |
+
---
|
| 255 |
+
|
| 256 |
+
## Repository layout
|
| 257 |
+
|
| 258 |
+
```
|
| 259 |
+
zerocg/ from-scratch NumPy framework (autograd, nn, AdamW, FLOP counter)
|
| 260 |
+
models/zeropoint.py ZeroPoint architecture + the small char-LM harness used in the experiments
|
| 261 |
+
models/transformer.py GPT-style Transformer baseline (learned or sinusoidal PE)
|
| 262 |
+
experiments/ compare / ablation / recall (+ recall_diag), length/ema/slot diagnostics
|
| 263 |
+
experiments/results/ dated archives of the reports (kept, annotated as historical)
|
| 264 |
+
tests/ numerical gradient checks, parallel=recurrent equivalence, golden reference
|
| 265 |
+
docs/ZeroPoint.md design document (Chinese; structured after "Attention Is All You Need")
|
| 266 |
+
docs/ZeroPoint_EN.md English translation of the design document
|
| 267 |
+
docs/AUDIT.md independent verification notes for the numbers in this card
|
| 268 |
+
```
|
| 269 |
+
|
| 270 |
+
## Reproducing
|
| 271 |
+
|
| 272 |
+
```bash
|
| 273 |
+
pip install numpy
|
| 274 |
+
python tests/test_zerocg.py # framework self-checks
|
| 275 |
+
python tests/test_zeropoint.py # architecture + parallel/recurrent equivalence (~20 s)
|
| 276 |
+
cd experiments
|
| 277 |
+
python compare.py # ~35 min on one CPU core
|
| 278 |
+
python ablation.py # ~70 min
|
| 279 |
+
python recall.py # ~15 min; known inconclusive (see above)
|
| 280 |
+
```
|
| 281 |
+
|
| 282 |
+
Experiments are pure NumPy on CPU. `compare.py` / `ablation.py` write a dated archive to
|
| 283 |
+
`experiments/results/` and update the latest report in `experiments/`.
|
| 284 |
+
|
| 285 |
+
---
|
| 286 |
+
|
| 287 |
+
## Data verification
|
| 288 |
+
|
| 289 |
+
The numbers in this card were independently audited, not merely copied from the project's
|
| 290 |
+
own reports:
|
| 291 |
+
|
| 292 |
+
- The **ZeroPoint side was re-derived bit-for-bit** from the saved diagnostic checkpoints
|
| 293 |
+
(`val64 = 1.5214 ± 0.0134`, per-seed 1.5311 / 1.5307 / 1.5024; the 32-sequence segment
|
| 294 |
+
values above; the ablation parameter counts). The Transformer checkpoints were never
|
| 295 |
+
saved, so those rows could not be re-derived (stated as a limitation above).
|
| 296 |
+
- Corrections applied after the audit: the 8-sequence length-extrapolation table was
|
| 297 |
+
superseded by the 32-sequence re-evaluation; the "effective write slots 3.2 / 5.0"
|
| 298 |
+
number was a metric bug (entropy over a non-normalised, gate-scaled write distribution —
|
| 299 |
+
the normalised values are ~6–10 write slots and ~10 read slots out of 16) and is not
|
| 300 |
+
used; "precise recall is a known weakness" was downgraded to "inconclusive"; "one matrix
|
| 301 |
+
multiply, as fast as attention" was corrected to "chunked parallel, one matmul per chunk";
|
| 302 |
+
"~110k parameters" was corrected to the exact 115,240.
|
| 303 |
+
- See [`docs/AUDIT.md`](docs/AUDIT.md) for the claim-by-claim audit table.
|
| 304 |
+
|
| 305 |
+
---
|
| 306 |
+
|
| 307 |
+
## Limitations (summary)
|
| 308 |
+
|
| 309 |
+
- Small scale: 115k params, 2 layers, 64-char context, 2,000 steps, one source-code corpus.
|
| 310 |
+
- Only validation loss was measured; no downstream tasks, no generation-quality evaluation,
|
| 311 |
+
no scaling study.
|
| 312 |
+
- The baseline is a minimal GPT-style decoder, not a tuned Transformer; the comparison is
|
| 313 |
+
fair in configuration but does not represent "Transformers" in general.
|
| 314 |
+
- 3 seeds; the LR was selected on seed 0 (small selection bias).
|
| 315 |
+
- No wall-clock or memory measurements; FLOP counts only.
|
| 316 |
+
- Length extrapolation: ZeroPoint degrades +0.037 nats/char at 4× training length; the
|
| 317 |
+
sinusoidal Transformer's collapse is a known property of absolute PE, measured with only
|
| 318 |
+
8 sequences.
|
| 319 |
+
- Exact recall: inconclusive; needs a corrected training configuration.
|
| 320 |
+
- The framework is a research prototype, not a production framework: no GPU kernels, no
|
| 321 |
+
distributed training, no mixed precision.
|
| 322 |
+
|
| 323 |
+
---
|
| 324 |
+
|
| 325 |
+
## Citation
|
| 326 |
+
|
| 327 |
+
```bibtex
|
| 328 |
+
@misc{zeropoint2026,
|
| 329 |
+
title = {ZeroPoint: a from-scratch NumPy framework and a hierarchical memory-slot
|
| 330 |
+
sequence architecture},
|
| 331 |
+
author = {Holmes2011},
|
| 332 |
+
year = {2026},
|
| 333 |
+
howpublished = {\url{https://huggingface.co/Holmes2011/ZeroPoint}},
|
| 334 |
+
note = {v0 research prototype; architecture and framework release, no pretrained weights}
|
| 335 |
+
}
|
| 336 |
+
```
|
| 337 |
+
|
| 338 |
+
## License
|
| 339 |
+
|
| 340 |
+
Apache-2.0 (code and documentation). This repository contains **no model weights, no
|
| 341 |
+
training corpus, and no training runs** of the project's separate 1.4M-parameter model.
|
| 342 |
+
|
| 343 |
+
---
|
| 344 |
+
|
| 345 |
+
## 中文摘要
|
| 346 |
+
|
| 347 |
+
**ZeroPoint 是一个全新的框架,不是已有模型的封装、微调或衍生品。** 本仓库从零实现了两部分:
|
| 348 |
+
|
| 349 |
+
1. **ZeroCG**:纯 Python + NumPy 的深度学习框架 —— 反向自动求导、`nn` 层、AdamW、矩阵乘
|
| 350 |
+
FLOP 统计,不依赖 PyTorch。
|
| 351 |
+
2. **ZeroPoint**:用一组**遗忘速度可学习、按存储层级初始化**的记忆槽取代自注意力。训练用
|
| 352 |
+
**分块并行**(块长 32,块间串行传状态),推理用**逐字递推**,状态固定、与文本长度无关,
|
| 353 |
+
**不需要位置编码**。
|
| 354 |
+
|
| 355 |
+
**对比条件(引用任何数字前请先看)**:Python 标准库源码 1,558,395 字符、字符级、词表 96;
|
| 356 |
+
三个模型同宽 64 / 2 层 / 4 头 / MLP 256 / Pre-LN / GELU;上下文 64、batch 32、2000 步、
|
| 357 |
+
AdamW、100 步预热 + 余弦、学习率在 {1e-3, 3e-3, 6e-3} 里按种子 0 选(三者都选 6e-3)、
|
| 358 |
+
种子 0/1/2;验证 loss = 40 个固定 batch × 32 序列 × 64 字的平均交叉熵;单机 CPU、NumPy。
|
| 359 |
+
|
| 360 |
+
**主要结果**:
|
| 361 |
+
|
| 362 |
+
- 头对头验证 loss:ZeroPoint **1.5214 ± 0.0134**,Transformer **1.5452 ± 0.0100**,
|
| 363 |
+
正弦位置编码 Transformer **1.5271 ± 0.0083**。最大差 0.0238 < 2× 种子标准差 0.0268
|
| 364 |
+
→ **统计上没有差别,不声称 ZeroPoint 更好**。
|
| 365 |
+
- 推理状态:ZeroPoint 固定 2,048 个浮点数(8 KiB fp32);Transformer KV 缓存在第 64 字时
|
| 366 |
+
16,384(8×)、第 1024 字时 262,144(**128×**)、第 4096 字时 1,048,576(**512×**)。
|
| 367 |
+
- 长度外推:训练 64 字,4 倍长度 loss 升高 **0.037**,64 倍长度升高 **0.219 ± 0.063**;
|
| 368 |
+
正弦位置编码的 Transformer 出训练长度后崩到约 4.64(≈随机水平;只有 8 条序列的测量)。
|
| 369 |
+
- 消融:分层半衰期初始化(+0.051)、可学习 λ(+0.028)、多头(1 头 +0.028)、去掉位置编码
|
| 370 |
+
(+0.061)有明确作用;槽数 4~32、写强度门、写入竞争在该规模下**无显著差异**。
|
| 371 |
+
- 精确回忆:**无结论**(学习率配置问题),不能用来比较架构。
|
| 372 |
+
|
| 373 |
+
**本仓库不含训练权重**:项目里另一个 1.4M 参数的模型(内部名 ZeroX)及其语料、训练运行
|
| 374 |
+
都不在本次发布范围内。
|
| 375 |
+
|
| 376 |
+
许可证:Apache-2.0。设计文档中文原文见 [`docs/ZeroPoint.md`](docs/ZeroPoint.md),
|
| 377 |
+
英文版见 [`docs/ZeroPoint_EN.md`](docs/ZeroPoint_EN.md)。
|
docs/AUDIT.md
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Independent audit of the numbers in this release
|
| 2 |
+
|
| 3 |
+
This document records how the numbers in `README.md` and `docs/ZeroPoint*.md` were
|
| 4 |
+
checked. The audit was done independently of the project's own summaries: every
|
| 5 |
+
publishable number was re-derived from the saved checkpoints where possible, or checked
|
| 6 |
+
arithmetically against the reports. **No training run was performed for the audit.**
|
| 7 |
+
|
| 8 |
+
## What was reproducible
|
| 9 |
+
|
| 10 |
+
Four trained ZeroPoint checkpoints from the comparison/diagnostic experiments were kept
|
| 11 |
+
(`experiments/out/diag_*` in the private source tree; not shipped here — they are not
|
| 12 |
+
needed, the experiment scripts retrain them in minutes). Loading them and re-running the
|
| 13 |
+
documented evaluation protocols reproduced, **bit-for-bit**:
|
| 14 |
+
|
| 15 |
+
| quantity | re-derived value |
|
| 16 |
+
|---|---|
|
| 17 |
+
| 64-char validation loss (3 seeds) | 1.5214 ± 0.0134 (per seed 1.5311 / 1.5307 / 1.5024) |
|
| 18 |
+
| 8-sequence length segments | 1.5550 / 1.6164 / 1.6410 / 1.7040 |
|
| 19 |
+
| 32-sequence length segments | 1.4372 / 1.4742 / 1.6617 / 1.6932; far−near +0.2190 ± 0.0630 |
|
| 20 |
+
| parameter counts (all ablation rows) | 115,240 base; 115,366 / 115,210 / 115,318 / 115,222 / 114,720 / 115,112 / 115,240 |
|
| 21 |
+
| inference state | 2,048 floats fixed; KV cache 16,384 / 262,144 / 1,048,576 at 64 / 1,024 / 4,096 |
|
| 22 |
+
| matmul FLOPs per char @ T=64 | ZeroPoint 232,448; Transformer 241,664 |
|
| 23 |
+
|
| 24 |
+
The Transformer baselines were **not** saved. Their reported rows (1.5452 ± 0.0100 and
|
| 25 |
+
1.5271 ± 0.0083) and the sinusoidal-PE extrapolation collapse (4.64) therefore rest on the
|
| 26 |
+
original report and could not be re-derived; they are labelled accordingly in the card.
|
| 27 |
+
|
| 28 |
+
## Claim-by-claim status
|
| 29 |
+
|
| 30 |
+
| claim | status | what the card says |
|
| 31 |
+
|---|---|---|
|
| 32 |
+
| ~110k parameters | corrected | exact 115,240 (Transformer 116,576; sin-PE 112,480) |
|
| 33 |
+
| validation loss "on par" | confirmed as a null result | 1.5214 vs 1.5452, diff 0.0238 < 2× std 0.0268 → not significant; no superiority claim |
|
| 34 |
+
| inference state 1/128 at the 1024th char | confirmed (with the position stated) | 2,048 vs 262,144 floats; 8× at 64, 512× at 4,096 |
|
| 35 |
+
| handles text 4× training length | confirmed, with degradation | +0.037 nats/char at 4×; +0.219 at 64× |
|
| 36 |
+
| training "one matmul over the whole sequence, as fast as attention" | corrected | chunked parallel (chunk 32, O(n/32) sequential steps); wall-clock not measured |
|
| 37 |
+
| forward compute 232K vs 242K FLOPs/char | confirmed (T=64 only) | stated together with the length dependence (Transformer 2,306,048 at 4,096) |
|
| 38 |
+
| length-extrapolation table (8 sequences) | superseded | replaced by 32-sequence values; old table kept only in the archived report with a note |
|
| 39 |
+
| Transformer(sin) collapse beyond training length | confirmed at 8 sequences | 4.64 vs 1.45, ≈ chance; N=8 caveat and "not evidence about attention in general" |
|
| 40 |
+
| "recurrent vs parallel algorithm" caused the 10-18 correction | corrected | the two forms agree to 1.4e-5 in float32; the real change was 8 → 32 validation sequences |
|
| 41 |
+
| "16 slots, only 3–5 used" (slot collapse) | corrected | metric bug (entropy on a non-normalised, gate-scaled write distribution); normalised: ~6–10 write / ~10 read slots out of 16, pooling-dependent — not used as a claim |
|
| 42 |
+
| ablation: hierarchy / learnable λ / no PE / heads | confirmed | +0.051 / +0.028 / +0.061 / +0.028, all above 2× max std; 8-head row marginal |
|
| 43 |
+
| ablation: 4/32 slots, write gate, write competition | confirmed as null | all below the decision line at this scale |
|
| 44 |
+
| precise recall is a "known weakness" | corrected | inconclusive: constant lr 3e-3 traps both architectures at loss ≈ ln 2; 1e-3 cosine solves N=2 |
|
| 45 |
+
| "32 switch combinations" in the self-check | corrected | 64 combinations (ema flag added later); parallel ≡ recurrent to < 1e-8 (float64) |
|
| 46 |
+
| exact recall / batch-size / calibration / long-run results | out of scope | those belong to the separate 1.4M model, which is not part of this release |
|
| 47 |
+
|
| 48 |
+
## Why no "beats the Transformer" claim
|
| 49 |
+
|
| 50 |
+
Under the project's pre-registered rule (a difference counts only if it exceeds 2× the
|
| 51 |
+
seed-to-seed standard deviation), the head-to-head gap is 0.89× the threshold. A standard-
|
| 52 |
+
error criterion would put the gap at ~2.5σ in ZeroPoint's favour, but that criterion was
|
| 53 |
+
not pre-registered, only 3 seeds were run, and the learning rate was selected on one of the
|
| 54 |
+
same seeds. The card therefore says "statistically indistinguishable" and reports the raw
|
| 55 |
+
numbers so readers can apply their own criterion.
|
| 56 |
+
|
| 57 |
+
## Practical checks that ship with the code
|
| 58 |
+
|
| 59 |
+
- `tests/test_zerocg.py` — numerical gradient checks for the framework's autograd.
|
| 60 |
+
- `tests/test_zeropoint.py` — all 64 switch combinations, parallel vs recurrent equivalence
|
| 61 |
+
(< 1e-8 float64), and a fixed golden reference in `tests/golden/`.
|
| 62 |
+
- `tests/make_golden.py` — regenerates the golden reference.
|
docs/ZeroPoint.md
ADDED
|
@@ -0,0 +1,247 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ZeroPoint:分层记忆槽网络
|
| 2 |
+
|
| 3 |
+
> 为 zeroOS 设计的序列模型架构。用它搭的语言模型叫 **ZeroX**,训练框架是 **ZeroCG**。(原名 ZeroMem。)文档结构参照《Attention Is All You Need》(Vaswani et al., 2017),
|
| 4 |
+
> 每节开头注明对应论文的哪一节,方便对照阅读。
|
| 5 |
+
>
|
| 6 |
+
> 代码:[`models/zeropoint.py`](../models/zeropoint.py) ·
|
| 7 |
+
> 自检:[`tests/test_zeropoint.py`](../tests/test_zeropoint.py) ·
|
| 8 |
+
> 对照实验:[`experiments/compare.py`](../experiments/compare.py) → [`compare.md`](../experiments/compare.md) ·
|
| 9 |
+
> 消融实验:[`experiments/ablation.py`](../experiments/ablation.py) → [`ablation.md`](../experiments/ablation.md) ·
|
| 10 |
+
> 精确回忆:[`experiments/recall.py`](../experiments/recall.py) → [`recall.md`](../experiments/recall.md) ·
|
| 11 |
+
> 历次结果存档:[`experiments/results/`](../experiments/results/)
|
| 12 |
+
>
|
| 13 |
+
> 状态:**v0(研究原型)**。
|
| 14 |
+
|
| 15 |
+
---
|
| 16 |
+
|
| 17 |
+
## 摘要
|
| 18 |
+
|
| 19 |
+
Transformer 用自注意力取代了循环,换来了训练时的完全并行;代价是推理时每生成一个字都要回看之前所有的字,
|
| 20 |
+
KV 缓存随文本长度线性增长。这对 GPU 服务器不是问题,对**没有 GPU、内存按页精打细算的操作系统内核**却是。
|
| 21 |
+
|
| 22 |
+
我们提出 **ZeroPoint**:用一组**带不同遗忘速度的记忆槽**取代自注意力。每读一个字,模型决定把它**写进**哪个槽、
|
| 23 |
+
从哪些槽**读出**;槽的遗忘速度按"寄存器 → 缓存 → 内存"的存储层级初始化,并在训练中学习。
|
| 24 |
+
ZeroPoint 有两种数学上等价的算法:训练时用**分块并行**(块长 32,每块一次矩阵乘,块间串行传递状态),
|
| 25 |
+
推理时**逐字递推**,状态大小固定,与文本长度无关,也不需要位置编码。
|
| 26 |
+
|
| 27 |
+
在 **115,240** 参数、1,558,395 字符的字符级语言建模上,ZeroPoint 与同规模 Transformer 的验证 loss
|
| 28 |
+
在 2× 种子标准差内无显著差异(1.5214 ± 0.0134 vs 1.5452 ± 0.0100,差 0.0238,判定线 0.0268;**不声称更好**)。
|
| 29 |
+
推理状态是后者在第 1024 个字时的 **1/128**(2,048 vs 262,144 个浮点数),
|
| 30 |
+
4 倍训练长度的文本上 loss 升高 0.037 nats/字符(64 倍时为 0.219)。
|
| 31 |
+
|
| 32 |
+
## 1 引言 〔论文 §1〕
|
| 33 |
+
|
| 34 |
+
论文 §1 的出发点是:循环网络"天生是串行的",第 t 步必须等第 t−1 步算完,没法在一条样本内并行。
|
| 35 |
+
自注意力把任意两个位置之间的路径缩短到 O(1),训练可以完全并行。
|
| 36 |
+
|
| 37 |
+
但论文没有讨论的另一面是**推理**:生成第 t 个字时,自注意力要和前面 t−1 个字的 Key/Value 逐一比较。
|
| 38 |
+
这些 Key/Value 必须一直存着(KV 缓存)。对 zeroOS 来说:
|
| 39 |
+
|
| 40 |
+
| | GPU 服务器 | zeroOS 内核 |
|
| 41 |
+
|---|---|---|
|
| 42 |
+
| 内存 | 几十 GB 显存 | 内核堆,按 4KB 页分配 |
|
| 43 |
+
| 算力 | 上万个并行单元 | 一个 CPU 核,x87 浮点 |
|
| 44 |
+
| 运行方式 | 一次处理一大批 | 用户敲一个字处理一个字 |
|
| 45 |
+
|
| 46 |
+
**我们想要的是两者的好处:训练几乎像 Transformer 一样并行(分块并行,串行步数是长度的 1/32),推理像 RNN 一样只带一个固定大小的状态。**
|
| 47 |
+
|
| 48 |
+
## 2 背景 〔论文 §2〕
|
| 49 |
+
|
| 50 |
+
论文 §2 提到 End-to-End Memory Networks [34] 用"循环的注意力"读外部记忆。ZeroPoint 和这条线以及
|
| 51 |
+
后来的"线性注意力 / 带衰减的循环模型"家族是近亲,不是凭空出现的:
|
| 52 |
+
|
| 53 |
+
| 已有工作 | 相同点 | ZeroPoint 的不同 |
|
| 54 |
+
|---|---|---|
|
| 55 |
+
| 记忆网络 / 神经图灵机 | 软寻址读写记忆槽 | 不需要逐步串行训练;每个槽有自己的遗忘速度 |
|
| 56 |
+
| 线性注意力(Katharopoulos 2020) | 并行 / 递推双形式 | 状态是 M 个可寻址的槽,不是 D×D 的外积矩阵 |
|
| 57 |
+
| RetNet(Sun 2023) | 衰减 λ^(t−u) 实现双形式、不需要位置编码 | λ 可学且按存储层级初始化;写入时 softmax 竞争选槽 |
|
| 58 |
+
| Mamba / 状态空间模型 | 递推推理、状态固定 | 不需要专门的扫描算子,一次矩阵乘就能并行 |
|
| 59 |
+
|
| 60 |
+
真正属于 ZeroPoint 的是这个组合:**多时间尺度的可学习遗忘 + 写时竞争选槽 + 写强度门 + 一次矩阵乘的并行形式**,
|
| 61 |
+
以及"为一个没有 GPU 的内核设计"这个出发点。
|
| 62 |
+
|
| 63 |
+
## 3 模型结构 〔论文 §3〕
|
| 64 |
+
|
| 65 |
+
### 3.1 层的堆叠 〔论文 §3.1〕
|
| 66 |
+
|
| 67 |
+
和论文的解码器一样:N 个相同的层堆叠,每层两个子层——**记忆混合层**(代替多头自注意力)和**逐位置前馈网络**,
|
| 68 |
+
每个子层外面套残差连接和 LayerNorm。
|
| 69 |
+
|
| 70 |
+
两处和论文不同:
|
| 71 |
+
|
| 72 |
+
- **LayerNorm 放在子层前面(Pre-LN)**:`x + Sublayer(LayerNorm(x))`。论文用的是 Post-LN
|
| 73 |
+
`LayerNorm(x + Sublayer(x))`,后来的 GPT-2 等工作发现 Pre-LN 训练更稳,不太依赖学习率预热。
|
| 74 |
+
对照组 Transformer 用的也是 Pre-LN,两边一致。
|
| 75 |
+
- **只有解码器(自回归)**。论文是做翻译的编码器-解码器;zeroOS 的任务(续写、意图识别)用不上编码器。
|
| 76 |
+
|
| 77 |
+
### 3.2 记忆混合层 〔论文 §3.2 注意力���
|
| 78 |
+
|
| 79 |
+
#### 3.2.1 读与写 〔对应 §3.2.1 缩放点积注意力〕
|
| 80 |
+
|
| 81 |
+
论文的注意力是"拿查询去和所有的键比"。ZeroPoint 反过来:**每个字在写入时就决定自己存到哪里**,
|
| 82 |
+
读的时候只和 M 个槽打交道,不和之前的字直接打交道。
|
| 83 |
+
|
| 84 |
+
设第 t 个字的向量为 `x_t`,记忆槽为 `S[1..M]`(每个是 D 维向量),保留率为 `λ[1..M]`:
|
| 85 |
+
|
| 86 |
+
```
|
| 87 |
+
写地址 w_t = softmax(x_t · Ww) × sigmoid(x_t · Wg) 写到哪些槽 × 写多重
|
| 88 |
+
写入值 v_t = x_t · Wv
|
| 89 |
+
更新 S[m] ← λ[m] · S[m] + w_t[m] · v_t 所有槽先按自己的速度遗忘,再写入
|
| 90 |
+
读地址 a_t = softmax(x_t · Wr) 从哪些槽读
|
| 91 |
+
读出 r_t = Σ_m a_t[m] · S[m]
|
| 92 |
+
输出 LayerNorm(r_t) · Wo
|
| 93 |
+
```
|
| 94 |
+
|
| 95 |
+
- **写地址用 softmax(槽之间竞争)**:一个字要么写进短期槽,要么写进长期槽,模型得决定"这个字值不值得长期记住"。
|
| 96 |
+
- **写强度门**:空格、换行这类没信息量的字可以几乎不写,不冲掉已有的记忆。
|
| 97 |
+
- **读出后做 LayerNorm**:长期槽可能累积了几十个字,短期槽只有一两个,量级差别很大。
|
| 98 |
+
这和论文 §3.2.1 除以 √d_k 的动机相同——别让数值落到 softmax / 后续层不好处理的区间。
|
| 99 |
+
- **槽的初始状态 S0 可学习**(不是全 0):否则序列开头几个字读出来的向量接近 0,
|
| 100 |
+
后面的 LayerNorm 会把噪声放大。
|
| 101 |
+
|
| 102 |
+
#### 3.2.2 多头 〔对应 §3.2.2 多头注意力〕
|
| 103 |
+
|
| 104 |
+
和论文一样分 H 个头,每头 D = d/H 维,各自有一套读写地址和 M 个槽,最后拼起来乘 Wo。
|
| 105 |
+
论文的理由同样适用:单头会把不同的"记什么"平均到一起。
|
| 106 |
+
|
| 107 |
+
#### 3.2.3 保留率与存储层级 〔ZeroPoint 独有〕
|
| 108 |
+
|
| 109 |
+
```
|
| 110 |
+
λ = 0.5 + 0.5 · sigmoid(θ) λ ∈ (0.5, 1),θ 可学
|
| 111 |
+
```
|
| 112 |
+
|
| 113 |
+
初始化时,每头的 M 个槽的**半衰期**(旧内容衰减到一半要经过几个字)按等比数列从 1.5 铺到 64:
|
| 114 |
+
|
| 115 |
+
```
|
| 116 |
+
槽 0 λ ≈ 0.63 半衰期 1.5 个字 ← 像寄存器:记"刚才那个字"
|
| 117 |
+
槽 1~14 λ 逐渐变大 ← 像缓存:记"这个词""这一句"
|
| 118 |
+
槽 15 λ ≈ 0.989 半衰期 64 个字 ← 像内存:记"这一段在讲什么"
|
| 119 |
+
```
|
| 120 |
+
|
| 121 |
+
论文 §3.5 的正弦位置编码"波长从 2π 到 10000·2π 成等比数列",ZeroPoint 的半衰期等比铺开是同一个思路:
|
| 122 |
+
用一组不同尺度的分量覆盖从很近到很远的距离。区别是论文把尺度编码进**输入**,ZeroPoint 把尺度编码进**记忆本身**。
|
| 123 |
+
|
| 124 |
+
### 3.3 逐位置前馈网络 〔论文 §3.3〕
|
| 125 |
+
|
| 126 |
+
与论文相同的结构 `FFN(x) = act(x·W1 + b1)·W2 + b2`,内层宽度 4d(论文 512 → 2048 也是 4 倍)。
|
| 127 |
+
激活函数用 GELU(论文是 ReLU;GPT 系列之后普遍用 GELU),对照组也用 GELU。
|
| 128 |
+
|
| 129 |
+
### 3.4 嵌入与输出 〔论文 §3.4〕
|
| 130 |
+
|
| 131 |
+
可学习的字嵌入,输出前 LayerNorm + 线性层 + softmax。没有采用论文的"输入/输出嵌入共享权重",对照组也没有。
|
| 132 |
+
|
| 133 |
+
### 3.5 位置信息 〔论文 §3.5〕
|
| 134 |
+
|
| 135 |
+
论文说:"模型里没有循环也没有卷积,所以必须注入位置信息。"
|
| 136 |
+
|
| 137 |
+
**ZeroPoint 不需要位置编码。** 它本质上是一个(可并行的)循环结构:第 t 个字读到的内容里,
|
| 138 |
+
u 个字之前写入的信息被乘上了 λ^(t−u)。不同的槽衰减速度不同,所以"某条信息在哪些槽里还剩多少"
|
| 139 |
+
本身就告诉了模型它是多久以前的。消融实验 (G) 专门检验了这一点。
|
| 140 |
+
|
| 141 |
+
## 4 为什么用记忆槽 〔论文 §4,表 1〕
|
| 142 |
+
|
| 143 |
+
论文表 1 比较了三个指标:每层计算量、最少串行步数、任意两位置之间的最长路径。
|
| 144 |
+
我们加上对内核最重要的两列:推理时每个字的计算量、推理时要保存的状态。
|
| 145 |
+
|
| 146 |
+
n = 序列长度,d = 宽度,M = 每头槽数(实验里 M = 16,d = 64)。
|
| 147 |
+
|
| 148 |
+
| 层类型 | 每层计算量(训练) | 串行步数 | 最长路径 | 推理:每字计算 | 推理:状态大小 |
|
| 149 |
+
|---|---|---|---|---|---|
|
| 150 |
+
| 自注意力 | O(n²·d) | O(1) | O(1) | O(n·d),越往后越慢 | O(n·d),随长度增长 |
|
| 151 |
+
| 循环网络 | O(n·d²) | O(n) | O(n) | O(d²) | O(d) |
|
| 152 |
+
| **ZeroPoint** | O(n·c·(M+d))(分块并行,c = 32) | O(n/c) | **O(1)**,但信号按 λ^(t−u) 衰减 | **O(d²)** | **O(M·d),固定** |
|
| 153 |
+
|
| 154 |
+
关键在于 ZeroPoint 有两种等价算法:
|
| 155 |
+
|
| 156 |
+
**并行形式(训练用)**。把递推展开:
|
| 157 |
+
|
| 158 |
+
```
|
| 159 |
+
r_t = Σ_{u ≤ t} A[t,u] · v_u, A[t,u] = Σ_m a_t[m] · w_u[m] · λ[m]^(t−u)
|
| 160 |
+
```
|
| 161 |
+
|
| 162 |
+
利用 `λ^(t−u) = λ^t · λ^(−u)`:
|
| 163 |
+
|
| 164 |
+
```
|
| 165 |
+
A = ( a ⊙ λ^t ) @ ( w ⊙ λ^(−u) )ᵀ 再把 u > t 的部分清零(因果)
|
| 166 |
+
```
|
| 167 |
+
|
| 168 |
+
只要一次矩阵乘,和注意力的 `Q·Kᵀ` 同样能整段并行。但 `λ^(−u)` 最大是 2^u,float32 到 u ≈ 120 就溢出。
|
| 169 |
+
|
| 170 |
+
所以实际用的是**分块并行**:序列切成每块 c = 32 个字,进入一块时槽的状态是 `S_in`:
|
| 171 |
+
|
| 172 |
+
```
|
| 173 |
+
块内 r_i = Σ_{j ≤ i} A[i,j] · v_j A 同上,但 i、j 是块内下标��λ^(−j) 最大 2^31
|
| 174 |
+
跨块 r_i += (a_i ⊙ λ^(i+1)) @ S_in
|
| 175 |
+
块尾 S_out = λ^c ⊙ S_in + (w ⊙ λ^(c−1−j))ᵀ @ v
|
| 176 |
+
```
|
| 177 |
+
|
| 178 |
+
块内是矩阵乘,块和块之间只传一个 M × d 的状态。训练时每层计算量 O(n·c·(M+d)),
|
| 179 |
+
**序列多长都行**(自检里 float32 跑 1000 个字不溢出,和递推形式一致)。
|
| 180 |
+
|
| 181 |
+
**递推形式(推理用)**。逐字更新 `S`,每层只存 M·d 个浮点数,跟文本多长无关。
|
| 182 |
+
|
| 183 |
+
两种形式结果一致:自检在**所有 64 种开关组合**下、300 个字(跨多个块)上逐位比对,float64 误差 < 1e-8。
|
| 184 |
+
另有一份固定的参考数据 `tests/golden/zeropoint_ref.npz`,以后任何实现(包括内核里的 C 版本)都拿它对答案。
|
| 185 |
+
**训练好的模型可以直接用递推形式放进内核**:不用 KV 缓存,每层固定一块 M × d 的状态,不需要动态分配。
|
| 186 |
+
|
| 187 |
+
关于"最长路径 O(1)"要说实话:u 处写入的信息确实一步就能被 t 处读到,但已经乘上了 λ^(t−u)。
|
| 188 |
+
注意力能"原样找回"任意远处的某个字,ZeroPoint 只能找回**还没忘干净、而且没被其他写入冲淡**的信息。
|
| 189 |
+
|
| 190 |
+
## 5 训练 〔论文 §5〕
|
| 191 |
+
|
| 192 |
+
| | 论文 | ZeroPoint 实验 |
|
| 193 |
+
|---|---|---|
|
| 194 |
+
| 数据 | WMT 2014 英德,450 万句对,BPE 37000 | Python 标准库源码 150 万字符,字符级,词表 96 |
|
| 195 |
+
| 批次 | 约 25000 个词元 | 32 × 64 = 2048 个字符 |
|
| 196 |
+
| 硬件 | 8 块 P100 | 一台 Mac 的 CPU(numpy,自研框架 zerocg) |
|
| 197 |
+
| 优化器 | Adam β=(0.9, 0.98),公式 3(预热 4000 步后按 1/√步数衰减) | AdamW β=(0.9, 0.999),权重衰减 0.01;预热 100 步后余弦退火到 10% |
|
| 198 |
+
| 正则化 | 残差 dropout 0.1、标签平滑 0.1 | 无(模型小、训练步数少,没观察到过拟合) |
|
| 199 |
+
| 步数 | 10 万步(12 小时) | 2000 步(每个模型约 2 分钟) |
|
| 200 |
+
|
| 201 |
+
ZeroPoint 和对照组用**完全相同**的训练设置、相同的随机种子(决定初始化和数据顺序)。
|
| 202 |
+
|
| 203 |
+
## 6 结果 〔论文 §6〕
|
| 204 |
+
|
| 205 |
+
### 6.1 对照实验 〔论文 §6.1,表 2〕
|
| 206 |
+
|
| 207 |
+
> 结果见 [`experiments/compare.md`](../experiments/compare.md)(`compare.py` 自动生成)。
|
| 208 |
+
|
| 209 |
+
### 6.2 模型变体(消融) 〔论文 §6.2,表 3〕
|
| 210 |
+
|
| 211 |
+
> 结果见 [`experiments/ablation.md`](../experiments/ablation.md)(`ablation.py` 自动生成)。
|
| 212 |
+
|
| 213 |
+
### 6.3 长度外推 〔论文 §3.5 的猜想〕
|
| 214 |
+
|
| 215 |
+
论文选正弦位置编码的理由是它"**可能**可以外推到比训练时更长的序列",但没有给出实验。
|
| 216 |
+
我们在 64 字上训练 ZeroPoint 和正弦位置编码的 Transformer,然后在 4,096 字的文本上分段测试
|
| 217 |
+
(32 条验证序列 × 3 个种子,两种模型用同一种分块并行算法)。
|
| 218 |
+
|
| 219 |
+
- ZeroPoint:64~255 字(训练长度的 4 倍)loss 为 1.4742 ± 0.0341,比训练长度内(0~63:1.4372)高 **+0.037**;
|
| 220 |
+
1024~4095 字(64 倍)为 1.6932 ± 0.0305,高 **+0.219 ± 0.063**。**外推会变差,但没有崩掉。**
|
| 221 |
+
- 正弦位置编码的 Transformer:出了训练长度就升到约 4.64(≈ ln 96 = 4.56,即随机水平)。
|
| 222 |
+
这是 8 条序列 × 3 个种子的测量,基线权重未保留,无法用 32 条序列复测。
|
| 223 |
+
|
| 224 |
+
> 早期 8 条序列的结果(1.5550 / 1.6164 / 1.6410 / 1.7040)抽样噪声大,已由上面的 32 条序列结果取代。
|
| 225 |
+
> 原始数据见 [`experiments/results/`](../experiments/results/)。
|
| 226 |
+
|
| 227 |
+
### 6.4 精确回忆 〔无结论,不能用来比较架构〕
|
| 228 |
+
|
| 229 |
+
> 结果见 [`experiments/recall.md`](../experiments/recall.md)。
|
| 230 |
+
|
| 231 |
+
`recall.py` 的训练配置(恒定 lr 3e-3)把两个架构都困在"两个候选值各猜一半"的局部解(loss ≈ ln 2)。
|
| 232 |
+
诊断显示换成 1e-3 余弦退火后,只训 N=2 的 Transformer 在 8,000 步内到 100%。在按修正后的学习率重跑之前,
|
| 233 |
+
这张 20,000 步的表只能说明训练配置有问题,**不能说明 ZeroPoint 或 Transformer 谁更擅长精确回忆**。
|
| 234 |
+
|
| 235 |
+
## 7 结论与局限 〔论文 §7〕
|
| 236 |
+
|
| 237 |
+
**局限**(都是实话):
|
| 238 |
+
|
| 239 |
+
- **规模很小**:只在约 11.5 万参数、150 万字符、单一语料上测过,结论不能外推到大模型。
|
| 240 |
+
- **精确回忆**:只有 M 个槽,信息会被压缩、混合。**推测**(未验证)它不擅长原样复制很久以前的内容;现有 recall 实验因学习率配置问题无结论。
|
| 241 |
+
- **没有编码器-解码器版本**:论文的核心应用是翻译;ZeroPoint 目前只做了解码器。
|
| 242 |
+
|
| 243 |
+
**下一步**:
|
| 244 |
+
|
| 245 |
+
1. 如果精确回忆确实弱:加一层"滑动窗口注意力"(只看最近 32 个字,状态仍然固定)做混合结构;
|
| 246 |
+
2. 在 zeroOS 的意图识别任务上训练;
|
| 247 |
+
3. 用 C 实现递推形式放进内核(x87 浮点)——这是 ZeroPoint 存在的原因。
|
docs/ZeroPoint_EN.md
ADDED
|
@@ -0,0 +1,214 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ZeroPoint: Hierarchical Memory-Slot Network
|
| 2 |
+
|
| 3 |
+
> A sequence-model architecture designed for zeroOS. The language model built on it is called **ZeroX**, and the training framework is **ZeroCG**. (Formerly named ZeroMem.) The document follows the structure of *Attention Is All You Need* (Vaswani et al., 2017); each section begins by noting which section of the paper it corresponds to, for easy cross-reading.
|
| 4 |
+
>
|
| 5 |
+
> Code: [`models/zeropoint.py`](../models/zeropoint.py) ·
|
| 6 |
+
> Self-checks: [`tests/test_zeropoint.py`](../tests/test_zeropoint.py) ·
|
| 7 |
+
> Comparison experiment: [`experiments/compare.py`](../experiments/compare.py) → [`compare.md`](../experiments/compare.md) ·
|
| 8 |
+
> Ablation experiment: [`experiments/ablation.py`](../experiments/ablation.py) → [`ablation.md`](../experiments/ablation.md) ·
|
| 9 |
+
> Exact recall: [`experiments/recall.py`](../experiments/recall.py) → [`recall.md`](../experiments/recall.md) ·
|
| 10 |
+
> Archived results: [`experiments/results/`](../experiments/results/)
|
| 11 |
+
>
|
| 12 |
+
> Status: **v0 (research prototype)**.
|
| 13 |
+
|
| 14 |
+
---
|
| 15 |
+
|
| 16 |
+
## Abstract
|
| 17 |
+
|
| 18 |
+
The Transformer replaced recurrence with self-attention, gaining full parallelism at training time; the price is that, at inference time, generating each character requires looking back at all previous characters, and the KV cache grows linearly with text length. This is not a problem for GPU servers, but it is for **an operating-system kernel with no GPU and memory meted out page by page**.
|
| 19 |
+
|
| 20 |
+
We propose **ZeroPoint**: replacing self-attention with a set of **memory slots that forget at different speeds**. As it reads each character, the model decides **which slot to write it into** and **from which slots to read**; the forgetting speeds of the slots are initialized according to a "register → cache → memory" storage hierarchy and are learned during training. ZeroPoint has two mathematically equivalent algorithms: at training time it uses a **chunked-parallel** form (chunk 32: one matrix multiplication per chunk, with a serial state carry between chunks), and at inference time it **recurs character by character**, with a state size that is fixed and independent of text length, and it needs no positional encoding.
|
| 21 |
+
|
| 22 |
+
On character-level language modeling with **115,240** parameters and 1,558,395 characters, ZeroPoint is statistically indistinguishable from a same-size Transformer (1.5214 ± 0.0134 vs 1.5452 ± 0.0100; difference 0.0238 < 2× seed std = 0.0268 — **we do not claim it is better**). Its inference state is **1/128** of the latter's at the 1024th character (2,048 vs 262,144 floats), and at 4× the training length the loss rises by 0.037 nats/character (0.219 at 64×).
|
| 23 |
+
|
| 24 |
+
## 1 Introduction 〔Paper §1〕
|
| 25 |
+
|
| 26 |
+
The starting point of the paper's §1 is: recurrent networks "are inherently sequential" — step t has to wait for step t−1 to finish computing, so there is no parallelism within a single sample. Self-attention shortens the path between any two positions to O(1), so training can be fully parallel.
|
| 27 |
+
|
| 28 |
+
But the other side, which the paper does not discuss, is **inference**: when generating the t-th character, self-attention has to compare against the Key/Value of the previous t−1 characters one by one. These Key/Values must be kept all along (the KV cache). For zeroOS:
|
| 29 |
+
|
| 30 |
+
| | GPU server | zeroOS kernel |
|
| 31 |
+
|---|---|---|
|
| 32 |
+
| Memory | tens of GB of VRAM | kernel heap, allocated in 4KB pages |
|
| 33 |
+
| Compute | tens of thousands of parallel units | one CPU core, x87 floating point |
|
| 34 |
+
| Execution mode | process a large batch at once | process one character as the user types it |
|
| 35 |
+
|
| 36 |
+
**What we want is the best of both: training almost as parallel as a Transformer (chunked parallel; sequential steps = length/32), inference carrying only a single fixed-size state like an RNN.**
|
| 37 |
+
|
| 38 |
+
## 2 Background 〔Paper §2〕
|
| 39 |
+
|
| 40 |
+
The paper's §2 mentions that End-to-End Memory Networks [34] read external memory with "recurrent attention". ZeroPoint is a close relative of this line of work and of the later "linear attention / recurrent models with decay" family; it did not appear out of thin air:
|
| 41 |
+
|
| 42 |
+
| Prior work | Similarity | How ZeroPoint differs |
|
| 43 |
+
|---|---|---|
|
| 44 |
+
| Memory networks / Neural Turing Machine | soft-addressable read and write of memory slots | no need for step-by-step serial training; each slot has its own forgetting speed |
|
| 45 |
+
| Linear attention (Katharopoulos 2020) | parallel / recurrent dual form | the state is M addressable slots, not a D×D outer-product matrix |
|
| 46 |
+
| RetNet (Sun 2023) | decay λ^(t−u) realizes the dual form and needs no positional encoding | λ is learnable and initialized by storage hierarchy; writes compete for slots via softmax |
|
| 47 |
+
| Mamba / state-space models | recurrent inference, fixed state | no specialized scan operator is needed; one matrix multiplication suffices to parallelize |
|
| 48 |
+
|
| 49 |
+
What genuinely belongs to ZeroPoint is this combination: **multi-timescale learnable forgetting + write-time competition among slots + a write-strength gate + a parallel form that uses one matrix multiplication**, together with the starting point of "designing for a kernel with no GPU".
|
| 50 |
+
|
| 51 |
+
## 3 Model Architecture 〔Paper §3〕
|
| 52 |
+
|
| 53 |
+
### 3.1 Stacking of Layers 〔Paper §3.1〕
|
| 54 |
+
|
| 55 |
+
Same as the paper's decoder: N identical layers are stacked, and each layer has two sub-layers — the **memory mixing layer** (replacing multi-head self-attention) and the **position-wise feed-forward network** — with a residual connection and LayerNorm around each sub-layer.
|
| 56 |
+
|
| 57 |
+
There are two differences from the paper:
|
| 58 |
+
|
| 59 |
+
- **LayerNorm is placed before the sub-layer (Pre-LN)**: `x + Sublayer(LayerNorm(x))`. The paper uses Post-LN `LayerNorm(x + Sublayer(x))`; later work such as GPT-2 found that Pre-LN trains more stably and does not depend much on learning-rate warmup. The control Transformer also uses Pre-LN, so both sides are consistent.
|
| 60 |
+
- **Decoder only (autoregressive)**. The paper is an encoder-decoder for translation; zeroOS's tasks (continuation, intent recognition) do not need an encoder.
|
| 61 |
+
|
| 62 |
+
### 3.2 Memory Mixing Layer 〔Paper §3.2 Attention〕
|
| 63 |
+
|
| 64 |
+
#### 3.2.1 Reading and Writing 〔Corresponds to §3.2.1 Scaled Dot-Product Attention〕
|
| 65 |
+
|
| 66 |
+
The paper's attention is "take a query and compare it with all the keys". ZeroPoint does the reverse: **each character decides where to store itself at write time**, and at read time it deals only with the M slots, never directly with previous characters.
|
| 67 |
+
|
| 68 |
+
Let the vector of the t-th character be `x_t`, the memory slots be `S[1..M]` (each a D-dimensional vector), and the retention rates be `λ[1..M]`:
|
| 69 |
+
|
| 70 |
+
```
|
| 71 |
+
write address w_t = softmax(x_t · Ww) × sigmoid(x_t · Wg) which slots to write × how strongly
|
| 72 |
+
write value v_t = x_t · Wv
|
| 73 |
+
update S[m] ← λ[m] · S[m] + w_t[m] · v_t all slots first forget at their own speed, then are written
|
| 74 |
+
read address a_t = softmax(x_t · Wr) from which slots to read
|
| 75 |
+
read out r_t = Σ_m a_t[m] · S[m]
|
| 76 |
+
output LayerNorm(r_t) · Wo
|
| 77 |
+
```
|
| 78 |
+
|
| 79 |
+
- **Write address uses softmax (competition among slots)**: a character either goes into a short-term slot or into a long-term slot; the model has to decide whether "this character is worth remembering long-term".
|
| 80 |
+
- **Write-strength gate**: uninformative characters such as spaces and newlines can be written almost not at all, without washing out existing memories.
|
| 81 |
+
- **LayerNorm after reading**: long-term slots may have accumulated dozens of characters while short-term slots have only one or two, so their magnitudes differ greatly. This is the same motivation as dividing by √d_k in the paper's §3.2.1 — do not let the values fall into a range that softmax / subsequent layers handle poorly.
|
| 82 |
+
- **The slots' initial state S0 is learnable** (not all zeros): otherwise the vectors read out at the beginning of the sequence are close to 0, and the subsequent LayerNorm would amplify the noise.
|
| 83 |
+
|
| 84 |
+
#### 3.2.2 Multiple Heads 〔Corresponds to §3.2.2 Multi-Head Attention〕
|
| 85 |
+
|
| 86 |
+
As in the paper, the model is split into H heads, each of dimension D = d/H, each with its own set of read/write addresses and M slots; finally they are concatenated and multiplied by Wo. The paper's reasoning applies equally: a single head would average the different "what to remember" together.
|
| 87 |
+
|
| 88 |
+
#### 3.2.3 Retention Rates and the Storage Hierarchy 〔ZeroPoint-specific〕
|
| 89 |
+
|
| 90 |
+
ZeroPoint's retention rates realize a hierarchical, learnable decay:
|
| 91 |
+
|
| 92 |
+
```
|
| 93 |
+
λ = 0.5 + 0.5 · sigmoid(θ) λ ∈ (0.5, 1), θ learnable
|
| 94 |
+
```
|
| 95 |
+
|
| 96 |
+
At initialization, the **half-lives** of the M slots of each head (how many characters it takes for old content to decay to half) are spread geometrically from 1.5 to 64:
|
| 97 |
+
|
| 98 |
+
```
|
| 99 |
+
slot 0 λ ≈ 0.63 half-life 1.5 characters ← like a register: remember "the character just now"
|
| 100 |
+
slots 1~14 λ gradually larger ← like a cache: remember "this word" "this sentence"
|
| 101 |
+
slot 15 λ ≈ 0.989 half-life 64 characters ← like memory: remember "what this passage is about"
|
| 102 |
+
```
|
| 103 |
+
|
| 104 |
+
The paper's §3.5 sinusoidal positional encoding has "wavelengths in a geometric progression from 2π to 10000·2π"; ZeroPoint's geometric spread of half-lives is the same idea: use a set of components at different scales to cover distances from very near to very far. The difference is that the paper encodes scale into the **input**, whereas ZeroPoint encodes scale into the **memory itself**.
|
| 105 |
+
|
| 106 |
+
### 3.3 Position-wise Feed-Forward Network 〔Paper §3.3〕
|
| 107 |
+
|
| 108 |
+
Same structure as the paper, `FFN(x) = act(x·W1 + b1)·W2 + b2`, with inner width 4d (the paper's 512 → 2048 is also 4x). The activation is GELU (the paper uses ReLU; after the GPT series, GELU is common), and the control group uses GELU too.
|
| 109 |
+
|
| 110 |
+
### 3.4 Embeddings and Output 〔Paper §3.4〕
|
| 111 |
+
|
| 112 |
+
Learnable character embeddings; before the output there is LayerNorm + a linear layer + softmax. We do not adopt the paper's "share weights between input and output embeddings", and neither does the control group.
|
| 113 |
+
|
| 114 |
+
### 3.5 Positional Information 〔Paper §3.5〕
|
| 115 |
+
|
| 116 |
+
The paper says: "the model contains no recurrence and no convolution, so positional information must be injected."
|
| 117 |
+
|
| 118 |
+
**ZeroPoint does not need positional encoding.** It is essentially a (parallelizable) recurrent structure: in what the t-th character reads, information written u characters earlier has been multiplied by λ^(t−u). Different slots decay at different speeds, so "how much of a piece of information is left in which slots" itself tells the model how long ago it was. Ablation experiment (G) specifically tests this.
|
| 119 |
+
|
| 120 |
+
## 4 Why Memory Slots 〔Paper §4, Table 1〕
|
| 121 |
+
|
| 122 |
+
The paper's Table 1 compares three metrics: computation per layer, minimum number of sequential steps, and maximum path length between any two positions. We add the two columns that matter most for a kernel: computation per character at inference, and state to keep at inference.
|
| 123 |
+
|
| 124 |
+
n = sequence length, d = width, M = number of slots per head (in the experiments M = 16, d = 64).
|
| 125 |
+
|
| 126 |
+
| Layer type | Computation per layer (training) | Sequential steps | Maximum path | Inference: computation per character | Inference: state size |
|
| 127 |
+
|---|---|---|---|---|---|
|
| 128 |
+
| Self-attention | O(n²·d) | O(1) | O(1) | O(n·d), slower toward the end | O(n·d), grows with length |
|
| 129 |
+
| Recurrent network | O(n·d²) | O(n) | O(n) | O(d²) | O(d) |
|
| 130 |
+
| **ZeroPoint** | O(n·c·(M+d)) (chunked parallel, c = 32) | O(n/c) | **O(1)**, but the signal decays as λ^(t−u) | **O(d²)** | **O(M·d), fixed** |
|
| 131 |
+
|
| 132 |
+
The key is that ZeroPoint has two equivalent algorithms:
|
| 133 |
+
|
| 134 |
+
**Parallel form (for training)**. Expand the recurrence:
|
| 135 |
+
|
| 136 |
+
```
|
| 137 |
+
r_t = Σ_{u ≤ t} A[t,u] · v_u, A[t,u] = Σ_m a_t[m] · w_u[m] · λ[m]^(t−u)
|
| 138 |
+
```
|
| 139 |
+
|
| 140 |
+
Using `λ^(t−u) = λ^t · λ^(−u)`:
|
| 141 |
+
|
| 142 |
+
```
|
| 143 |
+
A = ( a ⊙ λ^t ) @ ( w ⊙ λ^(−u) )ᵀ then zero out the part where u > t (causal)
|
| 144 |
+
```
|
| 145 |
+
|
| 146 |
+
Only one matrix multiplication is needed, and it can process the whole sequence in parallel just like attention's `Q·Kᵀ`. But `λ^(−u)` is at most 2^u, and float32 overflows at u ≈ 120.
|
| 147 |
+
|
| 148 |
+
So what is actually used is **chunked parallel**: the sequence is split into blocks of c = 32 characters; when entering a block, the slot state is `S_in`:
|
| 149 |
+
|
| 150 |
+
```
|
| 151 |
+
within block r_i = Σ_{j ≤ i} A[i,j] · v_j A as above, but i, j are within-block indices; λ^(−j) at most 2^31
|
| 152 |
+
across blocks r_i += (a_i ⊙ λ^(i+1)) @ S_in
|
| 153 |
+
block end S_out = λ^c ⊙ S_in + (w ⊙ λ^(c−1−j))ᵀ @ v
|
| 154 |
+
```
|
| 155 |
+
|
| 156 |
+
Within a block it is matrix multiplication, and only an M × d state is passed between blocks. At training time the computation per layer is O(n·c·(M+d)), and **any sequence length works** (in the self-check, float32 runs 1000 characters without overflow, matching the recurrent form).
|
| 157 |
+
|
| 158 |
+
**Recurrent form (for inference)**. Update `S` character by character; each layer stores only M·d floating-point numbers, regardless of how long the text is.
|
| 159 |
+
|
| 160 |
+
The two forms agree: the self-check compares them bit by bit on **all 64 switch combinations** over 300 characters (crossing several blocks), with float64 error < 1e-8. There is also a fixed reference file `tests/golden/zeropoint_ref.npz`; any future implementation (including a C version in the kernel) can check its answers against it. **A trained model can be put directly into the kernel in recurrent form**: no KV cache, a fixed M × d state per layer, no dynamic allocation.
|
| 161 |
+
|
| 162 |
+
About "maximum path O(1)", we must be honest: information written at u can indeed be read at t in one step, but it has already been multiplied by λ^(t−u). Attention can "retrieve verbatim" a character from arbitrarily far away; ZeroPoint can only retrieve information **that has not yet been forgotten and has not been washed out by other writes**.
|
| 163 |
+
|
| 164 |
+
## 5 Training 〔Paper §5〕
|
| 165 |
+
|
| 166 |
+
| | Paper | ZeroPoint experiments |
|
| 167 |
+
|---|---|---|
|
| 168 |
+
| Data | WMT 2014 English-German, 4.5M sentence pairs, BPE 37000 | Python standard library source, 1.5M characters, character-level, vocabulary 96 |
|
| 169 |
+
| Batch | about 25,000 tokens | 32 × 64 = 2048 characters |
|
| 170 |
+
| Hardware | 8 P100s | one Mac's CPU (numpy, the in-house framework zerocg) |
|
| 171 |
+
| Optimizer | Adam β=(0.9, 0.98), Eq. 3 (4000-step warmup then decay ∝ 1/√step) | AdamW β=(0.9, 0.999), weight decay 0.01; 100-step warmup then cosine decay to 10% |
|
| 172 |
+
| Regularization | residual dropout 0.1, label smoothing 0.1 | none (small model, few training steps, no overfitting observed) |
|
| 173 |
+
| Steps | 100K steps (12 hours) | 2000 steps (about 2 minutes per model) |
|
| 174 |
+
|
| 175 |
+
ZeroPoint and the control group use **exactly the same** training settings and the same random seeds (which determine initialization and data order).
|
| 176 |
+
|
| 177 |
+
## 6 Results 〔Paper §6〕
|
| 178 |
+
|
| 179 |
+
### 6.1 Comparison Experiment 〔Paper §6.1, Table 2〕
|
| 180 |
+
|
| 181 |
+
> Results: [`experiments/compare.md`](../experiments/compare.md) (auto-generated by `compare.py`).
|
| 182 |
+
|
| 183 |
+
### 6.2 Model Variants (Ablation) 〔Paper §6.2, Table 3〕
|
| 184 |
+
|
| 185 |
+
> Results: [`experiments/ablation.md`](../experiments/ablation.md) (auto-generated by `ablation.py`).
|
| 186 |
+
|
| 187 |
+
### 6.3 Length Extrapolation 〔Conjecture in Paper §3.5〕
|
| 188 |
+
|
| 189 |
+
The paper's reason for choosing sinusoidal positional encoding is that it "may allow the model to extrapolate to sequence lengths longer than the ones encountered during training", but no experiment was given. We train ZeroPoint and the sinusoidal-PE Transformer on 64 characters, then evaluate on 4,096-character text in segments (32 validation sequences × 3 seeds, same chunked-parallel algorithm for both models).
|
| 190 |
+
|
| 191 |
+
- ZeroPoint: loss at 64–255 characters (4× the training length) is 1.4742 ± 0.0341, i.e. **+0.037** above the in-training-length value (0–63: 1.4372); at 1,024–4,095 characters (64×) it is 1.6932 ± 0.0305, i.e. **+0.219 ± 0.063**. **Extrapolation degrades gracefully but does degrade.**
|
| 192 |
+
- The sinusoidal-PE Transformer jumps to about 4.64 once outside the training length (≈ ln 96 = 4.56, i.e. chance level). This is an 8-sequence × 3-seed measurement; the baseline checkpoints were not retained, so it could not be re-run with 32 sequences.
|
| 193 |
+
|
| 194 |
+
> The earlier 8-sequence values (1.5550 / 1.6164 / 1.6410 / 1.7040) are noise-dominated and are superseded by the 32-sequence numbers above. Raw data: [`experiments/results/`](../experiments/results/).
|
| 195 |
+
|
| 196 |
+
### 6.4 Exact Recall 〔inconclusive; not an architecture comparison〕
|
| 197 |
+
|
| 198 |
+
> Results: [`experiments/recall.md`](../experiments/recall.md).
|
| 199 |
+
|
| 200 |
+
The training configuration used by `recall.py` (constant lr 3e-3) traps both architectures at the "guess one of two candidates" local solution (loss ≈ ln 2). Diagnostics show that with a 1e-3 cosine schedule, the Transformer trained on N=2 alone reaches 100% in 8,000 steps. Until the experiment is re-run with the corrected learning rate, the 20,000-step table only shows that the *training setup* was wrong; **it says nothing about whether ZeroPoint or Transformer is better at exact recall**.
|
| 201 |
+
|
| 202 |
+
## 7 Conclusion and Limitations 〔Paper §7〕
|
| 203 |
+
|
| 204 |
+
**Limitations** (all honest):
|
| 205 |
+
|
| 206 |
+
- **Very small scale**: tested only on about 115K parameters, 1.5M characters, and a single corpus; the conclusions cannot be extrapolated to large models.
|
| 207 |
+
- **Exact recall**: with only M slots, information is compressed and mixed. It is a **hypothesis (not measured)** that this hurts verbatim copying of long-ago content; the existing recall experiment is inconclusive because of its learning-rate configuration.
|
| 208 |
+
- **No encoder-decoder version**: the paper's core application is translation; ZeroPoint currently implements only the decoder.
|
| 209 |
+
|
| 210 |
+
**Next steps**:
|
| 211 |
+
|
| 212 |
+
1. If exact recall is indeed weak: add a "sliding-window attention" layer (looking only at the most recent 32 characters, with the state still fixed) to make a hybrid structure;
|
| 213 |
+
2. Train on zeroOS's intent-recognition task;
|
| 214 |
+
3. Implement the recurrent form in C and put it in the kernel (x87 floating point) — this is why ZeroPoint exists.
|
experiments/ablation.md
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ZeroPoint 消融实验(仿论文表 3;历史存档:原名 ZeroMem)
|
| 2 |
+
|
| 3 |
+
每行只改一处,其余和 base 相同。2000 步,lr 0.003,种子 [0, 1, 2],其余训练设置同 [`compare.md`](compare.md)。
|
| 4 |
+
**注意**:本表所有行固定 lr 0.003;`compare.md` 用的是每个模型选出的最佳 lr 0.006(base loss 1.5500 vs 1.5214)。两张表只在各自内部可比,不能混用。
|
| 5 |
+
参数量和 base 差超过 1% 的变体,自动调整 MLP 宽度对齐参数量。
|
| 6 |
+
Δ = 该行验证 loss 均值 − base 均值(**正数表示变差**)。
|
| 7 |
+
|
| 8 |
+
> 判断标准:每行的阈值 = 2 × max(base 种子间标准差 0.0073, 该行标准差);例如 λ 固定行 0.0194、8 头行 0.0146。|Δ| 小于阈值的差别**不能算数**。
|
| 9 |
+
|
| 10 |
+
| | 改动 | 参数量 | 验证 loss | Δ | 结论 |
|
| 11 |
+
|---|---|---|---|---|---|
|
| 12 |
+
| base | ZeroPoint:4 头 × 16 槽,写强度门,写入竞争,分层初始化,λ 可学,可学初始状态,无位置编码 | 115,240 | 1.5500 ± 0.0073 | — | — |
|
| 13 |
+
| (A) | 1 个头(每头 64 维)(MLP 宽度调为 306 以对齐参数) | 115,366 | 1.5776 ± 0.0036 | +0.0276 | 变差:这个部件有用 |
|
| 14 |
+
| (A) | 8 个头(每头 8 维)(MLP 宽度调为 189 以对齐参数) | 115,210 | 1.5673 ± 0.0067 | +0.0173 | 变差:这个部件有用 |
|
| 15 |
+
| (B) | 每头 4 个槽(MLP 宽度调为 311 以对齐参数) | 115,318 | 1.5466 ± 0.0087 | -0.0034 | 差别不显著 |
|
| 16 |
+
| (B) | 每头 32 个槽(MLP 宽度调为 183 以对齐参数) | 115,222 | 1.5624 ± 0.0055 | +0.0124 | 差别不显著 |
|
| 17 |
+
| (C) | 去掉写强度门 | 114,720 | 1.5522 ± 0.0072 | +0.0022 | 差别不显著 |
|
| 18 |
+
| (D) | 写入不竞争:softmax 换成每槽独立的 sigmoid | 115,240 | 1.5487 ± 0.0066 | -0.0013 | 差别不显著 |
|
| 19 |
+
| (E) | 不分层:所有槽初始半衰期相同(约 9.8 个字) | 115,240 | 1.6010 ± 0.0079 | +0.0510 | 变差:这个部件有用 |
|
| 20 |
+
| (F) | λ 固定在初始值,不学习 | 115,112 | 1.5784 ± 0.0097 | +0.0284 | 变差:这个部件有用 |
|
| 21 |
+
| (G) | 加上论文 §3.5 的正弦位置编码 | 115,240 | 1.6106 ± 0.0078 | +0.0606 | 变差:这个部件有用 |
|
experiments/ablation.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ZeroPoint 消融实验,仿《Attention Is All You Need》表 3:每次只改一处,看那个部件到底有没有用。
|
| 2 |
+
|
| 3 |
+
python3 experiments/ablation.py
|
| 4 |
+
|
| 5 |
+
- 每组 3 个种子,报告 均值 ± 标准差
|
| 6 |
+
- 参数量对齐:某个改动让参数量和 base 差超过 1% 时(比如改头数、槽数),自动调整 MLP 宽度把参数量拉回来,
|
| 7 |
+
这样变好变差就不是"参数多了/少了"造成的
|
| 8 |
+
- 学习率固定 3e-3(所有变体相同)
|
| 9 |
+
结果:experiments/ablation.md(最新)+ experiments/results/<时间>-ablation.md(存档)
|
| 10 |
+
"""
|
| 11 |
+
import numpy as np
|
| 12 |
+
|
| 13 |
+
from common import CFG, ITERS, prepare_data, save_result, train
|
| 14 |
+
from models import ZeroX
|
| 15 |
+
|
| 16 |
+
SEEDS = [0, 1, 2]
|
| 17 |
+
LR = 3e-3
|
| 18 |
+
BASE = dict(n_slot=16, **CFG)
|
| 19 |
+
|
| 20 |
+
# (行号, 说明, 相对 base 的改动)
|
| 21 |
+
VARIANTS = [
|
| 22 |
+
("base", "ZeroPoint:4 头 × 16 槽,写强度门,写入竞争,分层初始化,λ 可学,可学初始状态,无位置编码", {}),
|
| 23 |
+
("(A)", "1 个头(每头 64 维)", dict(n_head=1)),
|
| 24 |
+
("(A)", "8 个头(每头 8 维)", dict(n_head=8)),
|
| 25 |
+
("(B)", "每头 4 个槽", dict(n_slot=4)),
|
| 26 |
+
("(B)", "每头 32 个槽", dict(n_slot=32)),
|
| 27 |
+
("(C)", "去掉写强度门", dict(gate=False)),
|
| 28 |
+
("(D)", "写入不竞争:softmax 换成每槽独立的 sigmoid", dict(write="sigmoid")),
|
| 29 |
+
("(E)", "不分层:所有槽初始半衰期相同(约 9.8 个字)", dict(init="uniform")),
|
| 30 |
+
("(F)", "λ 固定在初始值,不学习", dict(learn_lambda=False)),
|
| 31 |
+
("(G)", "加上论文 §3.5 的正弦位置编码", dict(pos="sin")),
|
| 32 |
+
]
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def matched(V, change, target):
|
| 36 |
+
"""返回 (参数量对齐后的配置, 是否调整过 MLP)。"""
|
| 37 |
+
cfg = {**BASE, **change}
|
| 38 |
+
n = ZeroX(V, **cfg).num_parameters()
|
| 39 |
+
if abs(n - target) / target <= 0.01:
|
| 40 |
+
return cfg, False
|
| 41 |
+
C, L = cfg["n_embd"], cfg["n_layer"]
|
| 42 |
+
hidden = 4 * C + round((target - n) / (L * (2 * C + 1))) # 每个隐藏单元每层 2C+1 个参数
|
| 43 |
+
return {**cfg, "mlp_hidden": hidden}, True
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def main():
|
| 47 |
+
text, V, train_data, val_data = prepare_data()
|
| 48 |
+
target = ZeroX(V, **BASE).num_parameters()
|
| 49 |
+
print(f"消融实验:{len(VARIANTS)} 组 × 种子 {SEEDS},每次 {ITERS} 步,lr {LR}", flush=True)
|
| 50 |
+
results = []
|
| 51 |
+
for row, desc, change in VARIANTS:
|
| 52 |
+
cfg, adjusted = matched(V, change, target)
|
| 53 |
+
vals = []
|
| 54 |
+
for seed in SEEDS:
|
| 55 |
+
np.random.seed(seed)
|
| 56 |
+
vals.append(train(f"{row} {desc[:10]}", ZeroX(V, **cfg), train_data, val_data, seed, LR))
|
| 57 |
+
n = ZeroX(V, **cfg).num_parameters()
|
| 58 |
+
results.append((row, desc + (f"(MLP 宽度调为 {cfg['mlp_hidden']} 以对齐参数)" if adjusted else ""),
|
| 59 |
+
n, np.mean(vals), np.std(vals), vals))
|
| 60 |
+
print(f" => {row} {desc}: {np.mean(vals):.4f} ± {np.std(vals):.4f}", flush=True)
|
| 61 |
+
|
| 62 |
+
base, base_std = results[0][3], results[0][4]
|
| 63 |
+
lines = ["# ZeroPoint 消融实验(仿论文表 3)\n",
|
| 64 |
+
f"每行只改一处,其余和 base 相同。{ITERS} 步,lr {LR},种子 {SEEDS},其余训练设置同 [`compare.md`](compare.md)。",
|
| 65 |
+
"参数量和 base 差超过 1% 的变体,自动调整 MLP 宽度对齐参数量。",
|
| 66 |
+
"Δ = 该行验证 loss 均值 − base 均值(**正数表示变差**)。\n",
|
| 67 |
+
f"> 判断标准:base 的种子间标准差是 {base_std:.4f}。|Δ| 小于约 2 倍标准差({2 * base_std:.4f})的差别**不能算数**。\n",
|
| 68 |
+
"| | 改动 | 参数量 | 验证 loss | Δ | 结论 |", "|---|---|---|---|---|---|"]
|
| 69 |
+
for row, desc, n, mean, std, _ in results:
|
| 70 |
+
if row == "base":
|
| 71 |
+
lines.append(f"| {row} | {desc} | {n:,} | {mean:.4f} ± {std:.4f} | — | — |")
|
| 72 |
+
continue
|
| 73 |
+
d = mean - base
|
| 74 |
+
noise = 2 * max(base_std, std)
|
| 75 |
+
if noise == 0:
|
| 76 |
+
verdict = "只有 1 个种子,无法判断"
|
| 77 |
+
elif abs(d) < noise:
|
| 78 |
+
verdict = "差别不显著"
|
| 79 |
+
else:
|
| 80 |
+
verdict = "变差:这个部件有用" if d > 0 else "变好:这个部件可能多余"
|
| 81 |
+
lines.append(f"| {row} | {desc} | {n:,} | {mean:.4f} ± {std:.4f} | {d:+.4f} | {verdict} |")
|
| 82 |
+
save_result("ablation", lines)
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
if __name__ == "__main__":
|
| 86 |
+
main()
|
experiments/common.py
ADDED
|
@@ -0,0 +1,119 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""各个实验脚本共用的部分:语料、批次、训练循环、计算量统计、结果存档。"""
|
| 2 |
+
import datetime
|
| 3 |
+
import glob
|
| 4 |
+
import os
|
| 5 |
+
import sys
|
| 6 |
+
import sysconfig
|
| 7 |
+
import time
|
| 8 |
+
|
| 9 |
+
import numpy as np
|
| 10 |
+
|
| 11 |
+
HERE = os.path.dirname(os.path.abspath(__file__))
|
| 12 |
+
sys.path.insert(0, os.path.join(HERE, ".."))
|
| 13 |
+
import zerocg as zc # noqa: E402
|
| 14 |
+
from models import ZeroX # noqa: E402
|
| 15 |
+
from zerocg import functional as F # noqa: E402
|
| 16 |
+
|
| 17 |
+
ITERS = int(os.environ.get("ITERS", 2000))
|
| 18 |
+
T, B = 64, 32
|
| 19 |
+
CFG = dict(n_embd=64, n_head=4, n_layer=2)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def load_corpus(path=None):
|
| 23 |
+
"""默认语料:Python 标准库源码(纯 ASCII 的 .py 文件,约 150 万字符)。"""
|
| 24 |
+
if path:
|
| 25 |
+
return open(path, encoding="utf-8").read()
|
| 26 |
+
parts, total = [], 0
|
| 27 |
+
for p in sorted(glob.glob(os.path.join(sysconfig.get_paths()["stdlib"], "*.py"))):
|
| 28 |
+
try:
|
| 29 |
+
s = open(p, encoding="utf-8").read()
|
| 30 |
+
except (UnicodeDecodeError, OSError):
|
| 31 |
+
continue
|
| 32 |
+
if s.isascii():
|
| 33 |
+
parts.append(s)
|
| 34 |
+
total += len(s)
|
| 35 |
+
if total > 1_500_000:
|
| 36 |
+
break
|
| 37 |
+
return "\n".join(parts)
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def prepare_data(path=None):
|
| 41 |
+
text = load_corpus(path)
|
| 42 |
+
chars = sorted(set(text))
|
| 43 |
+
stoi = {c: i for i, c in enumerate(chars)}
|
| 44 |
+
data = np.array([stoi[c] for c in text], dtype=np.int64)
|
| 45 |
+
split = int(0.9 * len(data))
|
| 46 |
+
return text, len(chars), data[:split], data[split:]
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def batches(data, n, length, seed, batch=B):
|
| 50 |
+
r = np.random.default_rng(seed)
|
| 51 |
+
for _ in range(n):
|
| 52 |
+
ix = r.integers(0, len(data) - length - 1, batch)
|
| 53 |
+
yield np.stack([data[i:i + length] for i in ix]), np.stack([data[i + 1:i + length + 1] for i in ix])
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def loss_of(logits, y):
|
| 57 |
+
return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1)).item()
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def train(name, model, train_data, val_data, seed, lr=3e-3, iters=None, seq_len=None, batch=None):
|
| 61 |
+
"""返回验证 loss(固定在 64 字上评估,便于各实验比较)。
|
| 62 |
+
学习率:100 步线性预热 + 余弦退火到 10%;AdamW,权重衰减 0.01,梯度裁剪 1.0。
|
| 63 |
+
seq_len / batch 默认 T / B;改训练长度时请同时改 batch,保持每步的字数不变。"""
|
| 64 |
+
iters = iters or ITERS
|
| 65 |
+
opt = zc.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
|
| 66 |
+
t0 = time.time()
|
| 67 |
+
for it, (x, y) in enumerate(batches(train_data, iters, seq_len or T, seed, batch or B)):
|
| 68 |
+
opt.lr = lr * min(1, (it + 1) / 100) * (0.1 + 0.9 * 0.5 * (1 + np.cos(np.pi * it / iters)))
|
| 69 |
+
logits = model(x)
|
| 70 |
+
loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1))
|
| 71 |
+
opt.zero_grad()
|
| 72 |
+
loss.backward()
|
| 73 |
+
zc.optim.clip_grad_norm(model.parameters(), 1.0)
|
| 74 |
+
opt.step()
|
| 75 |
+
if it % 500 == 0:
|
| 76 |
+
print(f" [{name} seed={seed} lr={lr:g}] step {it:5d} train {loss.item():.3f} "
|
| 77 |
+
f"{time.time() - t0:.0f}s", flush=True)
|
| 78 |
+
return evaluate(model, val_data)
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
def evaluate(model, val_data, n=40):
|
| 82 |
+
model.eval()
|
| 83 |
+
with zc.no_grad():
|
| 84 |
+
val = float(np.mean([loss_of(model(x), y) for x, y in batches(val_data, n, T, 999)]))
|
| 85 |
+
model.train()
|
| 86 |
+
return val
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def forward_flops_per_char(model, V):
|
| 90 |
+
"""前向一次、平均每个字多少次浮点运算(只算矩阵乘)。训练一步约是它的 3 倍。"""
|
| 91 |
+
x = np.zeros((1, T), dtype=np.int64)
|
| 92 |
+
with zc.no_grad(), zc.count_flops() as c:
|
| 93 |
+
model(x)
|
| 94 |
+
return c.flops / T
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def slot_usage(model):
|
| 98 |
+
"""每层"有效槽数":最近一次 forward 的平均读/写地址分布的 exp(熵),对头取平均。
|
| 99 |
+
16 个槽都用得一样多 = 16;全挤在一个槽 = 1。"""
|
| 100 |
+
out = []
|
| 101 |
+
for b in model.blocks:
|
| 102 |
+
read, write = b.mix.usage
|
| 103 |
+
eff = lambda p: float(np.mean(np.exp(-(p * np.log(p + 1e-12)).sum(-1)))) # noqa: E731
|
| 104 |
+
out.append((eff(read), eff(write)))
|
| 105 |
+
return out
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def save_result(name, lines):
|
| 109 |
+
"""结果写两份:experiments/results/<时间>-<name>.md(存档,不覆盖)和 experiments/<name>.md(最新)。"""
|
| 110 |
+
text = "\n".join(lines) + "\n"
|
| 111 |
+
stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
|
| 112 |
+
os.makedirs(os.path.join(HERE, "results"), exist_ok=True)
|
| 113 |
+
for path in (os.path.join(HERE, "results", f"{stamp}-{name}.md"), os.path.join(HERE, f"{name}.md")):
|
| 114 |
+
open(path, "w").write(text)
|
| 115 |
+
print("\n" + text)
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def is_zeropoint(model):
|
| 119 |
+
return isinstance(model, ZeroX)
|
experiments/compare.md
ADDED
|
@@ -0,0 +1,53 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ZeroPoint vs Transformer 对照实验
|
| 2 |
+
|
| 3 |
+
> **历史说明(2026-11-08 统一)**:本报告生成于 2026-10-11,当时架构名为 ZeroMem;除名称统一为 ZeroPoint 外,正文保留原样。两处已过时的内容在文末加了更正注:① 长度外推只用了 8 条验证序列(已被 32 条序列的复测取代);② "有效槽数"的写列用了未归一化的门控分布,数值偏低。
|
| 4 |
+
|
| 5 |
+
- 语料:Python 标准库源码,1,558,395 字符,词表 96;90% 训练 / 10% 验证
|
| 6 |
+
- 三个模型:宽 64、2 层、4 头、Pre-LN、GELU MLP 4 倍;ZeroPoint 每头 16 个槽、可学习初始状态、分块并行(块长 32)、无位置编码
|
| 7 |
+
- 训练:上下文 64、batch 32、2000 步、AdamW(权重衰减 0.01)、预热 100 步 + 余弦、梯度裁剪 1.0
|
| 8 |
+
- 学习率:每个模型在 [0.001, 0.003, 0.006] 里各试一次(种子 0),取最好的再跑种子 [1, 2]
|
| 9 |
+
- 验证 loss = 40 个固定验证 batch 的平均交叉熵(nats/字符,越低越好),均值 ± 标准差(3 个种子)
|
| 10 |
+
- 计算量 = 前向时每个字的矩阵乘浮点运算次数(不受机器负载影响;训练一步约 3 倍)
|
| 11 |
+
|
| 12 |
+
| 模型 | 参数量 | 最佳学习率 | 验证 loss | 前向计算量(每字) | 推理状态(第 64 字 / 第 4096 字,浮点数) |
|
| 13 |
+
|---|---|---|---|---|---|
|
| 14 |
+
| Transformer | 116,576 | 0.006 | 1.5452 ± 0.0100 | 242 K | 16,384 / 1,048,576(KV 缓存) |
|
| 15 |
+
| Transformer (sin) | 112,480 | 0.006 | 1.5271 ± 0.0083 | 242 K | 16,384 / 1,048,576(KV 缓存) |
|
| 16 |
+
| ZeroPoint | 115,240 | 0.006 | 1.5214 ± 0.0134 | 232 K | 2,048 / 2,048(固定) |
|
| 17 |
+
|
| 18 |
+
## 学习率搜索(种子 0 的验证 loss)
|
| 19 |
+
|
| 20 |
+
| 模型 | lr=0.001 | lr=0.003 | lr=0.006 |
|
| 21 |
+
|---|---|---|---|
|
| 22 |
+
| Transformer | 1.6908 | 1.5530 | 1.5520 |
|
| 23 |
+
| Transformer (sin) | 1.8006 | 1.5939 | 1.5374 |
|
| 24 |
+
| ZeroPoint | 1.6682 | 1.5433 | 1.5311 |
|
| 25 |
+
|
| 26 |
+
## 长度外推(训练上下文 64;8 条验证序列 × 3 个种子平均)
|
| 27 |
+
|
| 28 |
+
| 位置 | ZeroPoint | Transformer (sin) | Transformer(可学习位置编码) |
|
| 29 |
+
|---|---|---|---|
|
| 30 |
+
| 0 ~ 63(训练见过的长度) | 1.5550 | 1.4512 | 能跑 |
|
| 31 |
+
| 64 ~ 255 | 1.6164 | 4.6359 | 跑不了(只有 64 个位置) |
|
| 32 |
+
| 256 ~ 1023 | 1.6410 | 4.6057 | 跑不了(只有 64 个位置) |
|
| 33 |
+
| 1024 ~ 4095 | 1.7040 | —(注意力矩阵太大,没测) | 跑不了(只有 64 个位置) |
|
| 34 |
+
|
| 35 |
+
## ZeroPoint 槽的使用情况(有效槽数,满分 16)
|
| 36 |
+
|
| 37 |
+
有效槽数 = 平均读/写地址分布的 exp(熵):16 个槽用得一样多 = 16,全挤在一个槽 = 1。
|
| 38 |
+
|
| 39 |
+
| 层 | 读 | 写 |
|
| 40 |
+
|---|---|---|
|
| 41 |
+
| 0 | 10.1 | 3.2 |
|
| 42 |
+
| 1 | 9.9 | 5.0 |
|
| 43 |
+
|
| 44 |
+
---
|
| 45 |
+
|
| 46 |
+
## 更正(2026-11-08 追加,不改上表)
|
| 47 |
+
|
| 48 |
+
1. **长度外推**:上表只用了 8 条验证序列 × 3 个种子,抽样噪声大(同一模型 0~63 段:8 条序列 1.5550,32 条序列 1.4372)。
|
| 49 |
+
32 条序列的复测:0~63 **1.4372 ± 0.0187**、64~255 **1.4742 ± 0.0341**、256~1023 **1.6617 ± 0.0234**、
|
| 50 |
+
1024~4095 **1.6932 ± 0.0305**(远端 − 近端 = **+0.2190 ± 0.0630**)。详见 `diagnose.md` / `ema_longrun.md`。
|
| 51 |
+
Transformer(sin) 的 4.6359 / 4.6057 仍是 8 条序列的测量(基线权重未保留),≈ ln 96 = 4.56,即随机水平。
|
| 52 |
+
2. **有效槽数**:`slot_usage()` 的写列用的是未归一化的门控分布(每层质量只有 0.47 / 0.59),所以 3.2 / 5.0 偏低。
|
| 53 |
+
归一化后读列不变(10.1 / 9.9),写列约 6.0 / 10.4(每层);按段池化 8.6~8.9 / 16。不能据此说"16 个槽只用 3~5 个"。
|
experiments/compare.py
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ZeroPoint vs Transformer 对照实验(字符级语言模型)。
|
| 2 |
+
|
| 3 |
+
python3 experiments/compare.py # 默认语料:Python 标准库源码(离线可用)
|
| 4 |
+
python3 experiments/compare.py my.txt # 换成自己的文本
|
| 5 |
+
ITERS=2000 python3 experiments/compare.py
|
| 6 |
+
|
| 7 |
+
三个模型宽度、层数、头数、MLP 完全一样,唯一区别是"token 混合层 + 位置编码":
|
| 8 |
+
Transformer 自注意力 + 可学习位置编码
|
| 9 |
+
Transformer (sin) 自注意力 + 论文 §3.5 的正弦位置编码(理论上能外推到更长序列)
|
| 10 |
+
ZeroPoint 记忆槽,无位置编码
|
| 11 |
+
|
| 12 |
+
流程:
|
| 13 |
+
1. 学习率搜索:每个模型在 LRS 里各试一次(种子 0),取验证 loss 最低的
|
| 14 |
+
2. 用各自最好的学习率再跑种子 1、2,报告 3 个种子的 均值 ± 标准差
|
| 15 |
+
3. 长度外推:训练只见过 64 个字,测到 4096 个字
|
| 16 |
+
结果:experiments/compare.md(最新)+ experiments/results/<时间>-compare.md(存档)
|
| 17 |
+
"""
|
| 18 |
+
import sys
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
|
| 22 |
+
from common import (CFG, ITERS, T, batches, evaluate, forward_flops_per_char, is_zeropoint, prepare_data,
|
| 23 |
+
save_result, slot_usage, train)
|
| 24 |
+
from models import TransformerLM, ZeroX
|
| 25 |
+
import zerocg as zc
|
| 26 |
+
|
| 27 |
+
LRS = [1e-3, 3e-3, 6e-3]
|
| 28 |
+
SEEDS = [0, 1, 2]
|
| 29 |
+
MODELS = {
|
| 30 |
+
"Transformer": lambda V: TransformerLM(V, T, **CFG),
|
| 31 |
+
"Transformer (sin)": lambda V: TransformerLM(V, T, pos="sin", **CFG),
|
| 32 |
+
"ZeroPoint": lambda V: ZeroX(V, n_slot=16, **CFG),
|
| 33 |
+
}
|
| 34 |
+
# 长度外推的分段;Transformer 的注意力矩阵是 长度² 的,4096 字放不进内存,只测到 1024
|
| 35 |
+
SEGMENTS = [(0, 64), (64, 256), (256, 1024), (1024, 4096)]
|
| 36 |
+
LONG_MAX = {"Transformer (sin)": 1024, "ZeroPoint": 4096}
|
| 37 |
+
LONG_N = 8 # 测试序列条数
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def long_context_loss(model, val_data, max_len):
|
| 41 |
+
x, y = next(batches(val_data, 1, max_len, 7, batch=LONG_N))
|
| 42 |
+
with zc.no_grad():
|
| 43 |
+
logits = model(x).data
|
| 44 |
+
logp = logits - logits.max(-1, keepdims=True)
|
| 45 |
+
logp = logp - np.log(np.exp(logp).sum(-1, keepdims=True))
|
| 46 |
+
nll = -np.take_along_axis(logp, y[..., None], -1)[..., 0]
|
| 47 |
+
return np.array([nll[:, a:b].mean() if b <= max_len else np.nan for a, b in SEGMENTS])
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def main():
|
| 51 |
+
path = sys.argv[1] if len(sys.argv) > 1 else None
|
| 52 |
+
text, V, train_data, val_data = prepare_data(path)
|
| 53 |
+
print(f"语料 {len(text):,} 字符,词表 {V},每次训练 {ITERS} 步", flush=True)
|
| 54 |
+
|
| 55 |
+
sweep = {name: {} for name in MODELS} # name -> lr -> 验证 loss(种子 0)
|
| 56 |
+
finals = {name: [] for name in MODELS} # name -> [各种子的验证 loss]
|
| 57 |
+
long_loss = {name: [] for name in LONG_MAX}
|
| 58 |
+
usage = []
|
| 59 |
+
|
| 60 |
+
def run(name, seed, lr):
|
| 61 |
+
np.random.seed(seed)
|
| 62 |
+
model = MODELS[name](V)
|
| 63 |
+
val = train(name, model, train_data, val_data, seed, lr)
|
| 64 |
+
print(f" => {name} seed={seed} lr={lr:g}: val {val:.4f}", flush=True)
|
| 65 |
+
return model, val
|
| 66 |
+
|
| 67 |
+
for name in MODELS:
|
| 68 |
+
best = None
|
| 69 |
+
for lr in LRS:
|
| 70 |
+
model, val = run(name, 0, lr)
|
| 71 |
+
sweep[name][lr] = val
|
| 72 |
+
if best is None or val < best[1]:
|
| 73 |
+
best = (model, val, lr)
|
| 74 |
+
best_lr = best[2]
|
| 75 |
+
trained = [best[0]] + [run(name, s, best_lr)[0] for s in SEEDS[1:]]
|
| 76 |
+
finals[name] = [best[1]] + [None] * (len(SEEDS) - 1)
|
| 77 |
+
for k, model in enumerate(trained):
|
| 78 |
+
if k > 0:
|
| 79 |
+
finals[name][k] = evaluate(model, val_data)
|
| 80 |
+
if name in LONG_MAX:
|
| 81 |
+
long_loss[name].append(long_context_loss(model, val_data, LONG_MAX[name]))
|
| 82 |
+
if is_zeropoint(model):
|
| 83 |
+
usage.append(slot_usage(model))
|
| 84 |
+
|
| 85 |
+
# ---- 写结果 ----
|
| 86 |
+
m0 = {name: make(V) for name, make in MODELS.items()}
|
| 87 |
+
lines = ["# ZeroPoint vs Transformer 对照实验\n",
|
| 88 |
+
f"- 语料:{'Python 标准库源码' if not path else path},{len(text):,} 字符,词表 {V};90% 训练 / 10% 验证",
|
| 89 |
+
f"- 三个模型:宽 {CFG['n_embd']}、{CFG['n_layer']} 层、{CFG['n_head']} 头、Pre-LN、GELU MLP 4 倍;"
|
| 90 |
+
"ZeroPoint 每头 16 个槽、可学习初始状态、分块并行(块长 32)、无位置编码",
|
| 91 |
+
f"- 训练:上下文 {T}、batch 32、{ITERS} 步、AdamW(权重衰减 0.01)、预热 100 步 + 余弦、梯度裁剪 1.0",
|
| 92 |
+
f"- 学习率:每个模型在 {LRS} 里各试一次(种子 0),取最好的再跑种子 {SEEDS[1:]}",
|
| 93 |
+
"- 验证 loss = 40 个固定验证 batch 的平均交叉熵(nats/字符,越低越好),均值 ± 标准差(3 个种子)",
|
| 94 |
+
"- 计算量 = 前向时每个字的矩阵乘浮点运算次数(不受机器负载影响;训练一步约 3 倍)\n",
|
| 95 |
+
"| 模型 | 参数量 | 最佳学习率 | 验证 loss | 前向计算量(每字) | 推理状态(第 64 字 / 第 4096 字,浮点数) |",
|
| 96 |
+
"|---|---|---|---|---|---|"]
|
| 97 |
+
for name, vals in finals.items():
|
| 98 |
+
m = m0[name]
|
| 99 |
+
best_lr = min(sweep[name], key=sweep[name].get)
|
| 100 |
+
state = (f"{m.state_floats():,} / {m.state_floats():,}(固定)" if is_zeropoint(m)
|
| 101 |
+
else f"{m.state_floats(64):,} / {m.state_floats(4096):,}(KV 缓存)")
|
| 102 |
+
lines.append(f"| {name} | {m.num_parameters():,} | {best_lr:g} | {np.mean(vals):.4f} ± {np.std(vals):.4f} | "
|
| 103 |
+
f"{forward_flops_per_char(m, V) / 1e3:.0f} K | {state} |")
|
| 104 |
+
|
| 105 |
+
lines += ["", "## 学习率搜索(种子 0 的验证 loss)\n",
|
| 106 |
+
"| 模型 | " + " | ".join(f"lr={lr:g}" for lr in LRS) + " |",
|
| 107 |
+
"|---|" + "---|" * len(LRS)]
|
| 108 |
+
for name, d in sweep.items():
|
| 109 |
+
lines.append(f"| {name} | " + " | ".join(f"{d[lr]:.4f}" for lr in LRS) + " |")
|
| 110 |
+
|
| 111 |
+
lines += ["", f"## 长度外推(训练上下文 {T};{LONG_N} 条验证序列 × 3 个种子平均)\n",
|
| 112 |
+
"| 位置 | ZeroPoint | Transformer (sin) | Transformer(可学习位置编码) |", "|---|---|---|---|"]
|
| 113 |
+
lz, ls = np.nanmean(long_loss["ZeroPoint"], axis=0), np.nanmean(long_loss["Transformer (sin)"], axis=0)
|
| 114 |
+
for k, (a, b) in enumerate(SEGMENTS):
|
| 115 |
+
tag = "(训练见过的长度)" if b <= T else ""
|
| 116 |
+
s = "—(注意力矩阵太大,没测)" if np.isnan(ls[k]) else f"{ls[k]:.4f}"
|
| 117 |
+
lines.append(f"| {a} ~ {b - 1}{tag} | {lz[k]:.4f} | {s} | {'能跑' if b <= T else f'跑不了(只有 {T} 个位置)'} |")
|
| 118 |
+
|
| 119 |
+
u = np.mean(usage, axis=0) # (层, 2)
|
| 120 |
+
lines += ["", "## ZeroPoint 槽的使用情况(有效槽数,满分 16)\n",
|
| 121 |
+
"有效槽数 = 平均读/写地址分布的 exp(熵):16 个槽用得一样多 = 16,全挤在一个槽 = 1。\n",
|
| 122 |
+
"| 层 | 读 | 写 |", "|---|---|---|"]
|
| 123 |
+
lines += [f"| {i} | {r:.1f} | {w:.1f} |" for i, (r, w) in enumerate(u)]
|
| 124 |
+
save_result("compare", lines)
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
if __name__ == "__main__":
|
| 128 |
+
main()
|
experiments/diagnose.md
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 预注册实验:查三个问题
|
| 2 |
+
|
| 3 |
+
> 写于实验之前(2026-10-18)。跑完只在下面追加结果,不改这一节。
|
| 4 |
+
|
| 5 |
+
**背景**:ema_longrun.md 否定了"槽值无界累积"假设。无 ema 的模型在 1024~4095 字处 loss 比 64~255 字处高约 0.22;
|
| 6 |
+
有 ema 的高约 1.77。compare.md 里 16 个槽只有 3~5 个被写入(**更正**:3~5 是按未归一化的门控分布算的偏低值;归一化后读约 10/16,写约 6~10/层)。
|
| 7 |
+
|
| 8 |
+
## 假设和预测
|
| 9 |
+
|
| 10 |
+
**H2(训练长度)**:长序列变差是因为训练只见过 64 字;模型在更长的位置上遇到的槽状态分布是训练时没见过的。
|
| 11 |
+
- 预测:训练长度改为 256(每步字数不变、其余相同)后,1024~4095 与 64~255 的 loss 差 **缩小到 T64 的一半以下**。
|
| 12 |
+
- 不成立时会看到:差距没有缩小(< 2 × 标准差的变化)。
|
| 13 |
+
|
| 14 |
+
**H3(ema 为什么更差)**:ema 把长期槽的写入乘上 (1-λ),训练时(64 字)长期槽几乎不起作用,模型学会只用短期槽;
|
| 15 |
+
到了长序列,读地址仍指向短期槽,但短期槽里的内容和训练时的统计不同。
|
| 16 |
+
- 预测 ③a:ema 模型的**写地址**集中在半衰期短的槽上:半衰期 > 16 字的槽,写入占比 < T64 模型的一半。
|
| 17 |
+
- 预测 ③b:ema 模型学到的长期槽 λ 更接近 1(半衰期更长),或者长期槽的读占比在远处位置上升。
|
| 18 |
+
- 不成立时会看到:两组的写入分布没有明显差别。
|
| 19 |
+
|
| 20 |
+
**H4(槽塌缩)**:塌缩不是训练失败,而是写地址的 softmax 竞争让少数槽赢者通吃;在不同序列位置上"赢的槽"不同,
|
| 21 |
+
所以整体统计(compare.md 的 3~5 个有效槽,见上面对该度量的更正)低估了实际用到的槽数。
|
| 22 |
+
- 预测 ④:按位置分段(0~63 / 64~255 / 256~1023 / 1024~4095)分别统计有效写槽数,各段的并集 > 整体有效槽数 × 1.5;
|
| 23 |
+
同时各段的读/写分布之间有明显差别(任意两段的 JS 散度 > 0.1)。
|
| 24 |
+
- 不成立时会看到:各段的分布几乎相同,用的始终是同几个槽。
|
| 25 |
+
|
| 26 |
+
**判定**:每条按上面的数值阈值判,3 个种子均值;差别和 2 × 种子间标准差比。
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
## 结果(实验之后写)
|
| 30 |
+
|
| 31 |
+
### 验证 loss(64 字)和长度外推(32 条 × 4096 字;均值 ± 标准差,3 个种子)
|
| 32 |
+
|
| 33 |
+
| 条件 | 64 字验证 | 0~63 | 64~255 | 256~1023 | 1024~4095 | 远−近 差 |
|
| 34 |
+
|---|---|---|---|---|---|---|
|
| 35 |
+
| T64 | 1.5214 ± 0.0134 | 1.4372 ± 0.0187 | 1.4742 ± 0.0341 | 1.6617 ± 0.0234 | 1.6932 ± 0.0305 | +0.2190 ± 0.0630 |
|
| 36 |
+
| T64-ema | 1.5182 ± 0.0076 | 1.4308 ± 0.0039 | 1.6476 ± 0.0529 | 3.1084 ± 0.2157 | 3.4153 ± 0.2402 | +1.7677 ± 0.2188 |
|
| 37 |
+
| T256 | 1.5572 ± 0.0060 | 1.4741 ± 0.0117 | 1.4486 ± 0.0098 | 1.5098 ± 0.0022 | 1.5155 ± 0.0048 | +0.0669 ± 0.0051 |
|
| 38 |
+
|
| 39 |
+
### H3:写入落在半衰期 > 16 字的槽上的占比(0~63 段),和学到的最长半衰期
|
| 40 |
+
|
| 41 |
+
| 条件 | 长期槽写入占比 | 最长半衰期(字) | 远处读长期槽占比的变化 |
|
| 42 |
+
|---|---|---|---|
|
| 43 |
+
| T64 | 0.061 ± 0.005 | 109 | +0.009 |
|
| 44 |
+
| T64-ema | 0.014 ± 0.001 | 103 | +0.007 |
|
| 45 |
+
|
| 46 |
+
### H4:槽的使用(T64,写地址;各层各头平均)
|
| 47 |
+
|
| 48 |
+
| 段 | 有效写槽数 |
|
| 49 |
+
|---|---|
|
| 50 |
+
| 0~63 | 8.6 |
|
| 51 |
+
| 64~255 | 8.7 |
|
| 52 |
+
| 256~1023 | 8.5 |
|
| 53 |
+
| 1024~4095 | 8.9 |
|
| 54 |
+
| 整体 | 8.7 |
|
| 55 |
+
|
| 56 |
+
各段用到的槽的并集 ÷ 整体有效槽数 = 0.68;各段分布之间最大的 JS 散度 = 0.010
|
| 57 |
+
|
| 58 |
+
## 判定(按预注册的标准)
|
| 59 |
+
|
| 60 |
+
- H2 训练长度:远−近差 T64 +0.2190 → T256 +0.0669 → **成立**
|
| 61 |
+
- H3a ema 写入避开长期槽:占比 0.061 → 0.014 → **成立**
|
| 62 |
+
- H3b ema 的长期槽更长或远处更依赖长期槽 → **不成立**
|
| 63 |
+
- H4 槽塌缩是分段使用不同的槽:并集比 0.68,JS 0.010 → **不成立**
|
| 64 |
+
|
| 65 |
+
模型已保存在 experiments/out/diag_*.npz,可以事后再分析。
|
experiments/diagnose.py
ADDED
|
@@ -0,0 +1,216 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""预注册实验:查 ZeroPoint 的三个问题 —— 长序列变差的原因、ema 为什么更差、槽为什么塌缩。
|
| 2 |
+
|
| 3 |
+
python3 experiments/diagnose.py
|
| 4 |
+
|
| 5 |
+
设计在跑之前写定(PREREG),跑完只追加结果。训练好的模型保存在 experiments/out/,供事后再分析。
|
| 6 |
+
结果:experiments/diagnose.md + results/ 存档。
|
| 7 |
+
"""
|
| 8 |
+
import os
|
| 9 |
+
|
| 10 |
+
import numpy as np
|
| 11 |
+
|
| 12 |
+
from common import CFG, HERE, ITERS, T, B, batches, prepare_data, save_result, train
|
| 13 |
+
from models import ZeroX
|
| 14 |
+
import zerocg as zc
|
| 15 |
+
|
| 16 |
+
SEEDS = [0, 1, 2]
|
| 17 |
+
LR = 6e-3
|
| 18 |
+
LONG_T, LONG_N = 4096, 32
|
| 19 |
+
SEGMENTS = [(0, 64), (64, 256), (256, 1024), (1024, 4096)]
|
| 20 |
+
OUT = os.path.join(HERE, "out")
|
| 21 |
+
CONDITIONS = { # 名字 -> (ema, 训练长度, batch);每步的字数都是 2048
|
| 22 |
+
"T64": (False, 64, 32),
|
| 23 |
+
"T64-ema": (True, 64, 32),
|
| 24 |
+
"T256": (False, 256, 8),
|
| 25 |
+
}
|
| 26 |
+
|
| 27 |
+
PREREG = """# 预注册实验:查三个问题
|
| 28 |
+
|
| 29 |
+
> 写于实验之前(2026-10-18)。跑完只在下面追加结果,不改这一节。
|
| 30 |
+
|
| 31 |
+
**背景**:ema_longrun.md 否定了"槽值无界累积"假设。无 ema 的模型在 1024~4095 字处 loss 比 64~255 字处高约 0.22;
|
| 32 |
+
有 ema 的高约 1.77。compare.md 里 16 个槽只有 3~5 个被写入。
|
| 33 |
+
|
| 34 |
+
## 假设和预测
|
| 35 |
+
|
| 36 |
+
**H2(训练长度)**:长序列变差是因为训练只见过 64 字;模型在更长的位置上遇到的槽状态分布是训练时没见过的。
|
| 37 |
+
- 预测:训练长度改为 256(每步字数不变、其余相同)后,1024~4095 与 64~255 的 loss 差 **缩小到 T64 的一半以下**。
|
| 38 |
+
- 不成立时会看到:差距没有缩小(< 2 × 标准差的变化)。
|
| 39 |
+
|
| 40 |
+
**H3(ema 为什么更差)**:ema 把长期槽的写入乘上 (1-λ),训练时(64 字)长期槽几乎不起作用,模型学会只用短期槽;
|
| 41 |
+
到了长序列,读地址仍指向短期槽,但短期槽里的内容和训练时的统计不同。
|
| 42 |
+
- 预测 ③a:ema 模型的**写地址**集中在半衰期短的槽上:半衰期 > 16 字的槽,写入占比 < T64 模型的一半。
|
| 43 |
+
- 预测 ③b:ema 模型学到的长期槽 λ 更接近 1(半衰期更长),或者长期槽的读占比在远处位置上升。
|
| 44 |
+
- 不成立时会看到:两组的写入分布没有明显差别。
|
| 45 |
+
|
| 46 |
+
**H4(槽塌缩)**:塌缩不是训练失败,而是写地址的 softmax 竞争让少数槽赢者通吃;在不同序列位置上"赢的槽"不同,
|
| 47 |
+
所以整体统计(compare.md 的 3~5 个有效槽)低估了实际用到的槽数。
|
| 48 |
+
- 预测 ④:按位置分段(0~63 / 64~255 / 256~1023 / 1024~4095)分别统计有效写槽数,各段的并集 > 整体有效槽数 × 1.5;
|
| 49 |
+
同时各段的读/写分布之间有明显差别(任意两段的 JS 散度 > 0.1)。
|
| 50 |
+
- 不成立时会看到:各段的分布几乎相同,用的始终是同几个槽。
|
| 51 |
+
|
| 52 |
+
**判定**:每条按上面的数值阈值判,3 个种子均值;差别和 2 × 种子间标准差比。
|
| 53 |
+
"""
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def long_nll(model, val_data):
|
| 57 |
+
out = []
|
| 58 |
+
for x, y in batches(val_data, 4, LONG_T, 7, batch=LONG_N // 4):
|
| 59 |
+
with zc.no_grad():
|
| 60 |
+
logits = model(x).data
|
| 61 |
+
logp = logits - logits.max(-1, keepdims=True)
|
| 62 |
+
logp = logp - np.log(np.exp(logp).sum(-1, keepdims=True))
|
| 63 |
+
out.append(-np.take_along_axis(logp, y[..., None], -1)[..., 0])
|
| 64 |
+
nll = np.concatenate(out)
|
| 65 |
+
return np.array([nll[:, a:b].mean() for a, b in SEGMENTS])
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def address_stats(model, val_data):
|
| 69 |
+
"""递推跑 4096 字,按位置分段统计每层每头的平均读/写地址分布。返回 (层, 段, H, M) 两个数组。"""
|
| 70 |
+
x, _ = next(batches(val_data, 1, LONG_T, 7, batch=4))
|
| 71 |
+
states = model.init_state(4)
|
| 72 |
+
L = len(model.blocks)
|
| 73 |
+
H, M = model.blocks[0].mix.H, model.blocks[0].mix.M
|
| 74 |
+
reads = np.zeros((L, len(SEGMENTS), H, M))
|
| 75 |
+
writes = np.zeros((L, len(SEGMENTS), H, M))
|
| 76 |
+
for t in range(LONG_T):
|
| 77 |
+
_, states = model.step(x[:, t], states, t)
|
| 78 |
+
seg = next(i for i, (a, b) in enumerate(SEGMENTS) if a <= t < b)
|
| 79 |
+
for l, blk in enumerate(model.blocks):
|
| 80 |
+
a, w = blk.mix.last_addr
|
| 81 |
+
reads[l, seg] += a.mean(0)
|
| 82 |
+
writes[l, seg] += w.mean(0)
|
| 83 |
+
for i, (a, b) in enumerate(SEGMENTS):
|
| 84 |
+
reads[:, i] /= b - a
|
| 85 |
+
writes[:, i] /= b - a
|
| 86 |
+
return reads, writes
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def eff(p):
|
| 90 |
+
"""exp(熵):有效槽数。p 最后一维是槽。"""
|
| 91 |
+
p = p / p.sum(-1, keepdims=True)
|
| 92 |
+
return np.exp(-(p * np.log(p + 1e-12)).sum(-1))
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
def js(p, q):
|
| 96 |
+
p, q = p / p.sum(), q / q.sum()
|
| 97 |
+
m = (p + q) / 2
|
| 98 |
+
kl = lambda a, b: (a * np.log((a + 1e-12) / (b + 1e-12))).sum() # noqa: E731
|
| 99 |
+
return 0.5 * kl(p, m) + 0.5 * kl(q, m)
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def half_lives(model):
|
| 103 |
+
"""每层每头每槽的半衰期(字)。"""
|
| 104 |
+
return np.array([np.log(0.5) / np.log(b.mix.retention().data) for b in model.blocks])
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def main():
|
| 108 |
+
text, V, train_data, val_data = prepare_data()
|
| 109 |
+
os.makedirs(OUT, exist_ok=True)
|
| 110 |
+
print(PREREG, flush=True)
|
| 111 |
+
res = {}
|
| 112 |
+
for name, (ema, seq_len, batch) in CONDITIONS.items():
|
| 113 |
+
rows = []
|
| 114 |
+
for seed in SEEDS:
|
| 115 |
+
np.random.seed(seed)
|
| 116 |
+
model = ZeroX(V, n_slot=16, ema=ema, **CFG)
|
| 117 |
+
val = train(name, model, train_data, val_data, seed, LR, seq_len=seq_len, batch=batch)
|
| 118 |
+
model.eval()
|
| 119 |
+
model.save(os.path.join(OUT, f"diag_{name}_s{seed}.npz"))
|
| 120 |
+
ll = long_nll(model, val_data)
|
| 121 |
+
reads, writes = address_stats(model, val_data)
|
| 122 |
+
hl = half_lives(model)
|
| 123 |
+
rows.append(dict(val=val, long=ll, reads=reads, writes=writes, hl=hl))
|
| 124 |
+
print(f" => {name} seed={seed}: val {val:.4f}; long {np.round(ll, 4)}", flush=True)
|
| 125 |
+
res[name] = rows
|
| 126 |
+
|
| 127 |
+
def mean(name, key):
|
| 128 |
+
return np.mean([r[key] for r in res[name]], axis=0)
|
| 129 |
+
|
| 130 |
+
def std(name, key):
|
| 131 |
+
return np.std([r[key] for r in res[name]], axis=0)
|
| 132 |
+
|
| 133 |
+
lines = [PREREG, "\n## 结果(实验之后写)\n",
|
| 134 |
+
"### 验证 loss(64 字)和长度外推(32 条 × 4096 字;均值 ± 标准差,3 个种子)\n",
|
| 135 |
+
"| 条件 | 64 字验证 | " + " | ".join(f"{a}~{b - 1}" for a, b in SEGMENTS) + " | 远−近 差 |",
|
| 136 |
+
"|---|---|" + "---|" * (len(SEGMENTS) + 1)]
|
| 137 |
+
gaps = {}
|
| 138 |
+
for name in res:
|
| 139 |
+
lm, ls = mean(name, "long"), std(name, "long")
|
| 140 |
+
g = np.array([r["long"][3] - r["long"][1] for r in res[name]])
|
| 141 |
+
gaps[name] = g
|
| 142 |
+
lines.append(f"| {name} | {mean(name, 'val'):.4f} ± {std(name, 'val'):.4f} | " +
|
| 143 |
+
" | ".join(f"{m:.4f} ± {s:.4f}" for m, s in zip(lm, ls)) + f" | {g.mean():+.4f} ± {g.std():.4f} |")
|
| 144 |
+
|
| 145 |
+
# ---- H2 ----
|
| 146 |
+
g64, g256 = gaps["T64"], gaps["T256"]
|
| 147 |
+
# T64 的远−近差必须本来就是正的(确实变差了),H2 才有意义
|
| 148 |
+
h2 = g64.mean() > 0 and g256.mean() < 0.5 * g64.mean() and (g64.mean() - g256.mean()) > 2 * max(g64.std(), g256.std())
|
| 149 |
+
|
| 150 |
+
# ---- H3:写入在长期槽上的占比;学到的半衰期 ----
|
| 151 |
+
def long_slot_write_share(name):
|
| 152 |
+
out = []
|
| 153 |
+
for r in res[name]:
|
| 154 |
+
share = []
|
| 155 |
+
for l in range(r["hl"].shape[0]):
|
| 156 |
+
mask = r["hl"][l] > 16 # (H, M)
|
| 157 |
+
w = r["writes"][l, 0] # 训练长度内(0~63 段)的写分布 (H, M)
|
| 158 |
+
share.append((w * mask).sum() / w.sum())
|
| 159 |
+
out.append(np.mean(share))
|
| 160 |
+
return np.array(out)
|
| 161 |
+
s64, sema = long_slot_write_share("T64"), long_slot_write_share("T64-ema")
|
| 162 |
+
h3a = sema.mean() < 0.5 * s64.mean()
|
| 163 |
+
hl64, hlema = mean("T64", "hl"), mean("T64-ema", "hl")
|
| 164 |
+
far_read_long = {}
|
| 165 |
+
for name in ("T64", "T64-ema"):
|
| 166 |
+
vals = []
|
| 167 |
+
for r in res[name]:
|
| 168 |
+
v = []
|
| 169 |
+
for l in range(r["hl"].shape[0]):
|
| 170 |
+
mask = r["hl"][l] > 16
|
| 171 |
+
v.append((r["reads"][l, 3] * mask).sum() / r["reads"][l, 3].sum() -
|
| 172 |
+
(r["reads"][l, 0] * mask).sum() / r["reads"][l, 0].sum())
|
| 173 |
+
vals.append(np.mean(v))
|
| 174 |
+
far_read_long[name] = np.array(vals)
|
| 175 |
+
h3b = np.max(hlema) > np.max(hl64) * 1.5 or far_read_long["T64-ema"].mean() > far_read_long["T64"].mean() + 0.1
|
| 176 |
+
|
| 177 |
+
lines += ["", "### H3:写入落在半衰期 > 16 字的槽上的占比(0~63 段),和学到的最长半衰期\n",
|
| 178 |
+
"| 条件 | 长期槽写入占比 | 最长半衰期(字) | 远处读长期槽占比的变化 |", "|---|---|---|---|"]
|
| 179 |
+
for name in ("T64", "T64-ema"):
|
| 180 |
+
sh = long_slot_write_share(name)
|
| 181 |
+
lines.append(f"| {name} | {sh.mean():.3f} ± {sh.std():.3f} | {np.max(mean(name, 'hl')):.0f} | "
|
| 182 |
+
f"{far_read_long[name].mean():+.3f} |")
|
| 183 |
+
|
| 184 |
+
# ---- H4:分段的有效写槽数、并集、JS 散度 ----
|
| 185 |
+
lines += ["", "### H4:槽的使用(T64,写地址;各层各头平均)\n",
|
| 186 |
+
"| 段 | 有效写槽数 |", "|---|---|"]
|
| 187 |
+
union_ratio, js_min = [], []
|
| 188 |
+
for r in res["T64"]:
|
| 189 |
+
w = r["writes"] # (L, seg, H, M)
|
| 190 |
+
overall = eff(w.mean(1)) # (L, H)
|
| 191 |
+
per_seg = eff(w) # (L, seg, H)
|
| 192 |
+
used = (w / w.sum(-1, keepdims=True)) > 1.0 / w.shape[-1] # 用得比平均多的槽 (L, seg, H, M)
|
| 193 |
+
union = used.any(1).sum(-1) # (L, H)
|
| 194 |
+
union_ratio.append((union / overall).mean())
|
| 195 |
+
js_vals = [js(w[l, i, h], w[l, j, h]) for l in range(w.shape[0]) for h in range(w.shape[2])
|
| 196 |
+
for i in range(len(SEGMENTS)) for j in range(i + 1, len(SEGMENTS))]
|
| 197 |
+
js_min.append(np.max(js_vals))
|
| 198 |
+
w_mean = mean("T64", "writes")
|
| 199 |
+
for i, (a, b) in enumerate(SEGMENTS):
|
| 200 |
+
lines.append(f"| {a}~{b - 1} | {eff(w_mean[:, i]).mean():.1f} |")
|
| 201 |
+
lines.append(f"| 整体 | {eff(w_mean.mean(1)).mean():.1f} |")
|
| 202 |
+
union_ratio, js_min = np.array(union_ratio), np.array(js_min)
|
| 203 |
+
h4 = union_ratio.mean() > 1.5 and js_min.mean() > 0.1
|
| 204 |
+
lines += [f"\n各段用到的槽的并集 ÷ 整体有效槽数 = {union_ratio.mean():.2f};各段分布之间最大的 JS 散度 = {js_min.mean():.3f}"]
|
| 205 |
+
|
| 206 |
+
lines += ["", "## 判定(按预注册的标准)\n",
|
| 207 |
+
f"- H2 训练长度:远−近差 T64 {g64.mean():+.4f} → T256 {g256.mean():+.4f} → **{'成立' if h2 else '不成立'}**",
|
| 208 |
+
f"- H3a ema 写入避开长期槽:占比 {s64.mean():.3f} → {sema.mean():.3f} → **{'成立' if h3a else '不成立'}**",
|
| 209 |
+
f"- H3b ema 的长期槽更长或远处更依赖长期槽 → **{'成立' if h3b else '不成立'}**",
|
| 210 |
+
f"- H4 槽塌缩是分段使用不同的槽:并集比 {union_ratio.mean():.2f},JS {js_min.mean():.3f} → **{'成立' if h4 else '不成立'}**",
|
| 211 |
+
"", "模型已保存在 experiments/out/diag_*.npz,可以事后再分析。"]
|
| 212 |
+
save_result("diagnose", lines)
|
| 213 |
+
|
| 214 |
+
|
| 215 |
+
if __name__ == "__main__":
|
| 216 |
+
main()
|
experiments/ema_longrun.md
ADDED
|
@@ -0,0 +1,55 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 预注册实验:长序列变差与 ema 写入
|
| 2 |
+
|
| 3 |
+
> 写于实验之前(2026-10-18)。跑完后只在下面追加结果,不改这一节。
|
| 4 |
+
|
| 5 |
+
**观察**:compare.md 里 ZeroPoint(无 ema)在 1024~4095 字处的 loss(1.70)高于 64~255 字处(1.62)。
|
| 6 |
+
|
| 7 |
+
**假设 H1**:原因是长期槽(λ 接近 1)的数值随序列长度无界累积,超出训练时(64 字)见过的范围。
|
| 8 |
+
|
| 9 |
+
**预测**:
|
| 10 |
+
- ① 直接测:推理时槽值的最大绝对值,无 ema 的模型在 4000 字处明显大于 64 字处(比值 > 2);有 ema 的模型比值接近 1(< 1.5)。
|
| 11 |
+
- ② 间接测:开 ema 重新训练(其余完全相同),1024~4095 字处的 loss 不再明显高于 64~255 字处:差值 < 2 × 种子间标准差。
|
| 12 |
+
|
| 13 |
+
**H1 不成立时会看到**:① 无 ema 的槽值没有明显增长;或 ② 开 ema 后远处 loss 照样升高。
|
| 14 |
+
|
| 15 |
+
**判定**:① 和 ② 都成立 → 接受 H1,且 ema 是有效的修法;只有 ② 成立 → ema 有效但原因另有其他;
|
| 16 |
+
都不成立 → 放弃 H1,下一步查学到的 λ 值。另外要求 ema 不损害短距离:64 字处验证 loss 的差别 < 2 × 标准差。
|
| 17 |
+
|
| 18 |
+
**对照**:两组用同样的 3 个种子、同样 2000 步、学习率 0.006、同样语料;外推用 32 条 4096 字的验证序列,
|
| 19 |
+
两组都用分块并行算法算 loss。(10-18 的修正是把验证序列从 8 条增加到 32 条;分块并行与递推形式数学等价,float32 下 logit 最大差 1.4e-5,不是算法差异。)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
## 结果(实验之后写)
|
| 23 |
+
|
| 24 |
+
### 64 字处的验证 loss(均值 ± 标准差,3 个种子)
|
| 25 |
+
|
| 26 |
+
| 模型 | 验证 loss |
|
| 27 |
+
|---|---|
|
| 28 |
+
| no-ema | 1.5214 ± 0.0134 |
|
| 29 |
+
| ema | 1.5182 ± 0.0076 |
|
| 30 |
+
|
| 31 |
+
### 长度外推(32 条 × 4096 字,分块并行算法;均值 ± 标准差)
|
| 32 |
+
|
| 33 |
+
| 位置 | no-ema | ema |
|
| 34 |
+
|---|---|---|
|
| 35 |
+
| 0 ~ 63 | 1.4372 ± 0.0187 | 1.4308 ± 0.0039 |
|
| 36 |
+
| 64 ~ 255 | 1.4742 ± 0.0341 | 1.6476 ± 0.0529 |
|
| 37 |
+
| 256 ~ 1023 | 1.6617 ± 0.0234 | 3.1084 ± 0.2157 |
|
| 38 |
+
| 1024 ~ 4095 | 1.6932 ± 0.0305 | 3.4153 ± 0.2402 |
|
| 39 |
+
|
| 40 |
+
### 槽值的最大绝对值(递推推理,3 个种子平均)
|
| 41 |
+
|
| 42 |
+
| 位置 | no-ema | ema |
|
| 43 |
+
|---|---|---|
|
| 44 |
+
| 64 | 9.08 | 1.60 |
|
| 45 |
+
| 256 | 10.07 | 2.34 |
|
| 46 |
+
| 1024 | 9.88 | 1.51 |
|
| 47 |
+
| 4000 | 16.07 | 2.27 |
|
| 48 |
+
|
| 49 |
+
## 判定(按预注册的标准)
|
| 50 |
+
|
| 51 |
+
- 预测 ①(槽值比值 4000 字 / 64 字):无 ema 1.93,有 ema 1.44 → **不成立**
|
| 52 |
+
- 预测 ②(ema 模型 1024~4095 与 64~255 的 loss 差):+1.7677,噪声阈值 0.4805 → **不成立**(对照:无 ema 的差是 +0.2190)
|
| 53 |
+
- ema 不损害短距离:64 字处 loss 差 -0.0033 → **成立**
|
| 54 |
+
|
| 55 |
+
**结论**:① ② 都不成立:**放弃 H1**,下一步查学到的 λ 值
|
experiments/ema_longrun.py
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""预注册实验:ZeroPoint 长序列变差的原因是不是槽值无界累积?ema 写入能不能修好?
|
| 2 |
+
|
| 3 |
+
python3 experiments/ema_longrun.py
|
| 4 |
+
|
| 5 |
+
实验设计在跑之前写定(见 PREREG),结果出来后不改判定标准。结果:experiments/ema_longrun.md + results/ 存档。
|
| 6 |
+
"""
|
| 7 |
+
import numpy as np
|
| 8 |
+
|
| 9 |
+
from common import CFG, ITERS, T, batches, evaluate, prepare_data, save_result, train
|
| 10 |
+
from models import ZeroX
|
| 11 |
+
import zerocg as zc
|
| 12 |
+
|
| 13 |
+
SEEDS = [0, 1, 2]
|
| 14 |
+
LR = 6e-3 # compare.md 里 ZeroPoint 的最佳学习率
|
| 15 |
+
LONG_T = 4096
|
| 16 |
+
LONG_N = 32 # 外推测试序列条数(compare.py 只有 8 条,这里加到 32)
|
| 17 |
+
SEGMENTS = [(0, 64), (64, 256), (256, 1024), (1024, 4096)]
|
| 18 |
+
PROBE_POS = [64, 256, 1024, 4000] # 测槽值范围的位置
|
| 19 |
+
|
| 20 |
+
PREREG = """# 预注册实验:长序列变差与 ema 写入
|
| 21 |
+
|
| 22 |
+
> 写于实验之前(2026-10-18)。跑完后只在下面追加结果,不改这一节。
|
| 23 |
+
|
| 24 |
+
**观察**:compare.md 里 ZeroPoint(无 ema)在 1024~4095 字处的 loss(1.70)高于 64~255 字处(1.62)。
|
| 25 |
+
|
| 26 |
+
**假设 H1**:原因是长期槽(λ 接近 1)的数值随序列长度无界累积,超出训练时(64 字)见过的范围。
|
| 27 |
+
|
| 28 |
+
**预测**:
|
| 29 |
+
- ① 直接测:推理时槽值的最大绝对值,无 ema 的模型在 4000 字处明显大于 64 字处(比值 > 2);有 ema 的模型比值接近 1(< 1.5)。
|
| 30 |
+
- ② 间接测:开 ema 重新训练(其余完全相同),1024~4095 字处的 loss 不再明显高于 64~255 字处:差值 < 2 × 种子间标准差。
|
| 31 |
+
|
| 32 |
+
**H1 不成立时会看到**:① 无 ema 的槽值没有明显增长;或 ② 开 ema 后远处 loss 照样升高。
|
| 33 |
+
|
| 34 |
+
**判定**:① 和 ② 都成立 → 接受 H1,且 ema 是有效的修法;只有 ② 成立 → ema 有效但原因另有其他;
|
| 35 |
+
都不成立 → 放弃 H1,下一步查学到的 λ 值。另外要求 ema 不损害短距离:64 字处验证 loss 的差别 < 2 × 标准差。
|
| 36 |
+
|
| 37 |
+
**对照**:两组用同样的 3 个种子、同样 {iters} 步、学习率 {lr}、同样语料;外推用 {n} 条 {lt} 字的验证序列,
|
| 38 |
+
两组都用分块并行算法算 loss(compare.py 里 ZeroPoint 用的是递推算法,和 Transformer 不对等,这里统一)。
|
| 39 |
+
"""
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def long_loss(model, val_data):
|
| 43 |
+
"""(各段 loss 均值) ;分 4 批算,省内存。"""
|
| 44 |
+
nll_all = []
|
| 45 |
+
for x, y in batches(val_data, 4, LONG_T, 7, batch=LONG_N // 4):
|
| 46 |
+
with zc.no_grad():
|
| 47 |
+
logits = model(x).data
|
| 48 |
+
logp = logits - logits.max(-1, keepdims=True)
|
| 49 |
+
logp = logp - np.log(np.exp(logp).sum(-1, keepdims=True))
|
| 50 |
+
nll_all.append(-np.take_along_axis(logp, y[..., None], -1)[..., 0])
|
| 51 |
+
nll = np.concatenate(nll_all)
|
| 52 |
+
return np.array([nll[:, a:b].mean() for a, b in SEGMENTS])
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def slot_magnitude(model, val_data):
|
| 56 |
+
"""递推跑 4096 字,记录各位置处所有层槽值的最大绝对值。返回 {位置: 最大|S|}。"""
|
| 57 |
+
x, _ = next(batches(val_data, 1, LONG_T, 7, batch=4))
|
| 58 |
+
states = model.init_state(4)
|
| 59 |
+
out = {}
|
| 60 |
+
for t in range(LONG_T):
|
| 61 |
+
_, states = model.step(x[:, t], states, t)
|
| 62 |
+
if t + 1 in PROBE_POS:
|
| 63 |
+
out[t + 1] = max(float(np.abs(S).max()) for S in states)
|
| 64 |
+
return out
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def main():
|
| 68 |
+
text, V, train_data, val_data = prepare_data()
|
| 69 |
+
print(PREREG.format(iters=ITERS, lr=LR, n=LONG_N, lt=LONG_T), flush=True)
|
| 70 |
+
res = {}
|
| 71 |
+
for ema in (False, True):
|
| 72 |
+
name = "ema" if ema else "no-ema"
|
| 73 |
+
vals, longs, mags = [], [], []
|
| 74 |
+
for seed in SEEDS:
|
| 75 |
+
np.random.seed(seed)
|
| 76 |
+
model = ZeroX(V, n_slot=16, ema=ema, **CFG)
|
| 77 |
+
vals.append(train(name, model, train_data, val_data, seed, LR))
|
| 78 |
+
model.eval()
|
| 79 |
+
longs.append(long_loss(model, val_data))
|
| 80 |
+
mags.append(slot_magnitude(model, val_data))
|
| 81 |
+
print(f" => {name} seed={seed}: val {vals[-1]:.4f}; long {np.round(longs[-1], 4)}; "
|
| 82 |
+
f"|S| {[round(mags[-1][p], 2) for p in PROBE_POS]}", flush=True)
|
| 83 |
+
res[name] = dict(val=np.array(vals), long=np.array(longs),
|
| 84 |
+
mag=np.array([[m[p] for p in PROBE_POS] for m in mags]))
|
| 85 |
+
|
| 86 |
+
# ---- 结果 + 按预注册的标准判定 ----
|
| 87 |
+
lines = [PREREG.format(iters=ITERS, lr=LR, n=LONG_N, lt=LONG_T), "\n## 结果(实验之后写)\n",
|
| 88 |
+
"### 64 字处的验证 loss(均值 ± 标准差,3 个种子)\n", "| 模型 | 验证 loss |", "|---|---|"]
|
| 89 |
+
for name, r in res.items():
|
| 90 |
+
lines.append(f"| {name} | {r['val'].mean():.4f} ± {r['val'].std():.4f} |")
|
| 91 |
+
lines += ["", f"### 长度外推({LONG_N} 条 × {LONG_T} 字,分块并行算法;均值 ± 标准差)\n",
|
| 92 |
+
"| 位置 | " + " | ".join(res) + " |", "|---|" + "---|" * len(res)]
|
| 93 |
+
for k, (a, b) in enumerate(SEGMENTS):
|
| 94 |
+
lines.append(f"| {a} ~ {b - 1} | " + " | ".join(f"{r['long'][:, k].mean():.4f} ± {r['long'][:, k].std():.4f}"
|
| 95 |
+
for r in res.values()) + " |")
|
| 96 |
+
lines += ["", "### 槽值的最大绝对值(递推推理,3 个种子平均)\n",
|
| 97 |
+
"| 位置 | " + " | ".join(res) + " |", "|---|" + "---|" * len(res)]
|
| 98 |
+
for k, p in enumerate(PROBE_POS):
|
| 99 |
+
lines.append(f"| {p} | " + " | ".join(f"{r['mag'][:, k].mean():.2f}" for r in res.values()) + " |")
|
| 100 |
+
|
| 101 |
+
no, em = res["no-ema"], res["ema"]
|
| 102 |
+
ratio_no = (no["mag"][:, 3] / no["mag"][:, 0]).mean()
|
| 103 |
+
ratio_em = (em["mag"][:, 3] / em["mag"][:, 0]).mean()
|
| 104 |
+
p1 = ratio_no > 2 and ratio_em < 1.5
|
| 105 |
+
far, near = em["long"][:, 3], em["long"][:, 1]
|
| 106 |
+
gap, noise = (far - near).mean(), 2 * max(far.std(), near.std())
|
| 107 |
+
p2 = gap < noise
|
| 108 |
+
gap_no = (no["long"][:, 3] - no["long"][:, 1]).mean()
|
| 109 |
+
short_ok = abs(em["val"].mean() - no["val"].mean()) < 2 * max(em["val"].std(), no["val"].std())
|
| 110 |
+
if p1 and p2:
|
| 111 |
+
verdict = "① ② 都成立:**接受 H1**,ema 是有效的修法"
|
| 112 |
+
elif p2:
|
| 113 |
+
verdict = "只有 ② 成立:ema 有效,但长序列变差的原因不是(或不只是)槽值无界"
|
| 114 |
+
elif p1:
|
| 115 |
+
verdict = "只有 ① 成立:槽值确实无界,但 ema 没有修好远处的 loss,H1 不足以解释变差"
|
| 116 |
+
else:
|
| 117 |
+
verdict = "① ② 都不成立:**放弃 H1**,下一步查学到的 λ 值"
|
| 118 |
+
lines += ["", "## 判定(按预注册的标准)\n",
|
| 119 |
+
f"- 预测 ①(槽值比值 4000 字 / 64 字):无 ema {ratio_no:.2f},有 ema {ratio_em:.2f} → **{'成立' if p1 else '不成立'}**",
|
| 120 |
+
f"- 预测 ②(ema 模型 1024~4095 与 64~255 的 loss 差):{gap:+.4f},噪声阈值 {noise:.4f} → **{'成立' if p2 else '不成立'}**"
|
| 121 |
+
f"(对照:无 ema 的差是 {gap_no:+.4f})",
|
| 122 |
+
f"- ema 不损害短距离:64 字处 loss 差 {em['val'].mean() - no['val'].mean():+.4f} → **{'成立' if short_ok else '不成立'}**",
|
| 123 |
+
f"\n**结论**:{verdict}" + ("" if short_ok else ";但 ema 损害了短距离效果,不能直接采用")]
|
| 124 |
+
save_result("ema_longrun", lines)
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
if __name__ == "__main__":
|
| 128 |
+
main()
|
experiments/recall.md
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 长距离精确回忆(联想检索)
|
| 2 |
+
|
| 3 |
+
`k1 v1 k2 v2 … kN vN [问] kq` → 回答 kq 对应的值。16 个键、16 个值,随机猜的准确率是 1/16 ≈ 6%。
|
| 4 |
+
三个模型宽 64、2 层、4 头;训练 20000 步,N 在 2~16 随机;每格是 512 道测试题 × 种子 [0, 1] 的平均准确率。
|
| 5 |
+
|
| 6 |
+
| 模型 | N = 2 | N = 4 | N = 8 | N = 12 | N = 16 |
|
| 7 |
+
|---|---|---|---|---|---|
|
| 8 |
+
| Transformer | 50% | 30% | 23% | 20% | 18% |
|
| 9 |
+
| ZeroPoint (16 槽) | 51% | 32% | 23% | 20% | 16% |
|
| 10 |
+
| ZeroPoint (32 槽) | 49% | 31% | 24% | 20% | 17% |
|
| 11 |
+
|
| 12 |
+
## 判定与解释(11-08 补)
|
| 13 |
+
|
| 14 |
+
**预注册判定**(见 [`recall_prereg.md`](recall_prereg.md))
|
| 15 |
+
- P1(Transformer 在 N=16 准确率 ≥ 80%)**不成立**:18%。
|
| 16 |
+
- P2(ZeroPoint 比 Transformer 低 ≥ 20 个百分点)**不成立**:N=16 时 16% / 17% vs 18%,差 1~2 点。
|
| 17 |
+
- P3(32 槽比 16 槽高 ≥ 5 个百分点)**不成立**:17% vs 16%。
|
| 18 |
+
|
| 19 |
+
**和 3000 步比**:几乎完全一样(Transformer N=2 52% → 50%、N=16 14% → 18%;ZeroPoint 16 槽 51% → 51%、17% → 16%)。训练量加长 6.7 倍没有任何可辨别的提升。
|
| 20 |
+
|
| 21 |
+
## 诊断(11-08,`recall_diag.py`;诊断性实验,不参与上面的预注册判定)
|
| 22 |
+
|
| 23 |
+
1. **固定 32 条 N=2 题反复训练**:600 步训练集准确率 100%,loss ≈ 0 → 框架、优化、反向传播没问题。
|
| 24 |
+
2. **只训 N=1**(复制值,不需要匹配查询):500 步 100% → 复制机制能学会。
|
| 25 |
+
3. **只训 N=2,换学习率安排**(种子 0,各 8000 步):
|
| 26 |
+
- 恒定 3e-3(本脚本原设置):N=2 停在 53%,loss ≈ 0.72。
|
| 27 |
+
- 余弦 3e-3 → 3e-4:N=2 停在 53~59%,loss ≈ 0.72。
|
| 28 |
+
- **余弦 1e-3 → 1e-4:第 6000 步开始突破,第 8000 步 N=2 = 100%,loss 0.001。**
|
| 29 |
+
→ 任务能学会;**原来的 3e-3 恒定学习率把模型困在"两个候选值各猜一半"的状态**(loss ≈ ln 2 ≈ 0.69)。另外,单独只用恒定 3e-3 训 N=2 到 20000 步也仍是 53%,所以不是步数问题。
|
| 30 |
+
4. **混合 N(2~16)用同一个 1e-3 余弦跑 12000 步**:N=2 仍只有 56%。混合训练里 N=2 的样本只占约 1/15,梯度信号被稀释,12000 步对 N=2 只相当于 800 步左右——还不够。
|
| 31 |
+
|
| 32 |
+
**结论(已测量)**:20000 步无结论不是"训练量差一点",而是**训练配置问题**:
|
| 33 |
+
- 恒定 3e-3 太高,模型停在 50/50 对冲的局部解,学不出"用查询键选值"的回路;
|
| 34 |
+
- 混合 N 把简单样本(N 小)的信号稀释了大约 15 倍。
|
| 35 |
+
- 只学 N=2 时,1e-3 余弦 8000 步就能到 100%,说明 Transformer 架构有能力学这个任务。
|
| 36 |
+
|
| 37 |
+
**下一步**:按修正后的学习率(1e-3 余弦)重跑 recall,或者加课程(先 N=2,再逐步加长到 16);需要重写预注册。在那之前,recall 的结论仍是"无结论",**不能拿这张 20000 步的表去比架构**。
|
experiments/recall.py
ADDED
|
@@ -0,0 +1,87 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""长距离精确回忆(联想检索)实验:专门测 ZeroPoint 的已知弱点。
|
| 2 |
+
|
| 3 |
+
序列格式(每个数字是一个词元):
|
| 4 |
+
k1 v1 k2 v2 ... kN vN [问] kq → 模型要说出 kq 对应的 v
|
| 5 |
+
键从 16 个里不重复地抽 N 个,值从另外 16 个里随机抽。要答对,模型必须把 N 对"键 → 值"原样记住。
|
| 6 |
+
注意力天生擅长这个(用 kq 当查询去找那个键);ZeroPoint 只有 M 个槽,对数一多就会混。
|
| 7 |
+
|
| 8 |
+
python3 experiments/recall.py
|
| 9 |
+
训练时 N 在 2~16 之间随机,测试时分别报告 N = 2, 4, 8, 12, 16 的准确率。
|
| 10 |
+
结果:experiments/recall.md(最新)+ experiments/results/<时间>-recall.md(存档)
|
| 11 |
+
"""
|
| 12 |
+
import os
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
|
| 16 |
+
from common import CFG, save_result
|
| 17 |
+
from models import TransformerLM, ZeroX
|
| 18 |
+
import zerocg as zc
|
| 19 |
+
from zerocg import functional as F
|
| 20 |
+
|
| 21 |
+
ITERS = int(os.environ.get("RECALL_ITERS", 3000))
|
| 22 |
+
NK, NV = 16, 16
|
| 23 |
+
QUERY = NK + NV # "问" 这个词元
|
| 24 |
+
VOCAB = NK + NV + 1
|
| 25 |
+
MAX_N = 16
|
| 26 |
+
MAX_LEN = 2 * MAX_N + 2
|
| 27 |
+
TEST_N = [2, 4, 8, 12, 16]
|
| 28 |
+
SEEDS = [0, 1]
|
| 29 |
+
MODELS = {
|
| 30 |
+
"Transformer": lambda: TransformerLM(VOCAB, MAX_LEN, **CFG),
|
| 31 |
+
"ZeroPoint (16 槽)": lambda: ZeroX(VOCAB, n_slot=16, **CFG),
|
| 32 |
+
"ZeroPoint (32 槽)": lambda: ZeroX(VOCAB, n_slot=32, **CFG),
|
| 33 |
+
}
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def make_batch(rng, B, n):
|
| 37 |
+
keys = np.stack([rng.permutation(NK)[:n] for _ in range(B)])
|
| 38 |
+
vals = rng.integers(NK, NK + NV, size=(B, n))
|
| 39 |
+
seq = np.empty((B, 2 * n + 2), dtype=np.int64)
|
| 40 |
+
seq[:, 0:2 * n:2], seq[:, 1:2 * n:2] = keys, vals
|
| 41 |
+
q = rng.integers(0, n, B)
|
| 42 |
+
seq[:, -2], seq[:, -1] = QUERY, keys[np.arange(B), q]
|
| 43 |
+
return seq, vals[np.arange(B), q]
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def train_and_eval(make, seed):
|
| 47 |
+
np.random.seed(seed)
|
| 48 |
+
rng = np.random.default_rng(seed)
|
| 49 |
+
model = make()
|
| 50 |
+
opt = zc.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
|
| 51 |
+
for it in range(ITERS):
|
| 52 |
+
opt.lr = 3e-3 * min(1, (it + 1) / 100)
|
| 53 |
+
x, y = make_batch(rng, 32, int(rng.integers(2, MAX_N + 1)))
|
| 54 |
+
loss = F.cross_entropy(model(x)[:, -1], y)
|
| 55 |
+
opt.zero_grad()
|
| 56 |
+
loss.backward()
|
| 57 |
+
zc.optim.clip_grad_norm(model.parameters(), 1.0)
|
| 58 |
+
opt.step()
|
| 59 |
+
model.eval()
|
| 60 |
+
test = np.random.default_rng(1000 + seed)
|
| 61 |
+
acc = {}
|
| 62 |
+
with zc.no_grad():
|
| 63 |
+
for n in TEST_N:
|
| 64 |
+
x, y = make_batch(test, 512, n)
|
| 65 |
+
acc[n] = float((model(x).data[:, -1].argmax(-1) == y).mean())
|
| 66 |
+
return acc
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def main():
|
| 70 |
+
print(f"联想检索:{len(MODELS)} 个模型 × 种子 {SEEDS},每次 {ITERS} 步", flush=True)
|
| 71 |
+
res = {}
|
| 72 |
+
for name, make in MODELS.items():
|
| 73 |
+
runs = [train_and_eval(make, s) for s in SEEDS]
|
| 74 |
+
res[name] = {n: np.mean([r[n] for r in runs]) for n in TEST_N}
|
| 75 |
+
print(f" => {name}: " + " ".join(f"N={n}:{res[name][n]:.0%}" for n in TEST_N), flush=True)
|
| 76 |
+
lines = ["# 长距离精确回忆(联想检索)\n",
|
| 77 |
+
"`k1 v1 k2 v2 … kN vN [问] kq` → 回答 kq 对应的值。16 个键、16 个值,随机猜的准确率是 1/16 ≈ 6%。",
|
| 78 |
+
f"三个模型宽 {CFG['n_embd']}、{CFG['n_layer']} 层、{CFG['n_head']} 头;训练 {ITERS} 步,N 在 2~{MAX_N} 随机;"
|
| 79 |
+
f"每格是 512 道测试题 × 种子 {SEEDS} 的平均准确率。\n",
|
| 80 |
+
"| 模型 | " + " | ".join(f"N = {n}" for n in TEST_N) + " |", "|---|" + "---|" * len(TEST_N)]
|
| 81 |
+
for name, acc in res.items():
|
| 82 |
+
lines.append(f"| {name} | " + " | ".join(f"{acc[n]:.0%}" for n in TEST_N) + " |")
|
| 83 |
+
save_result("recall", lines)
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
if __name__ == "__main__":
|
| 87 |
+
main()
|
experiments/recall_diag.py
ADDED
|
@@ -0,0 +1,118 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""精确回忆实验的诊断脚本(诊断性,不参与预注册判定)。
|
| 2 |
+
|
| 3 |
+
回答的问题:为什么 recall.py 跑 3000 步和 20000 步结果几乎一样(Transformer N=16 只有 14~18%)?
|
| 4 |
+
|
| 5 |
+
python3 experiments/recall_diag.py
|
| 6 |
+
|
| 7 |
+
四段,都对固定测试题评估、种子 0:
|
| 8 |
+
A 固定 32 条 N=2 题反复训练:能不能背下来(检验框架/优化)。
|
| 9 |
+
B 只训 N=1(复制值,不需要匹配):能不能学会复制。
|
| 10 |
+
C 只训 N=2,比较学习率安排:恒定 3e-3 / 余弦 3e-3→3e-4 / 余弦 1e-3→1e-4,各 8000 步。
|
| 11 |
+
D 混合 N(2~16)用 1e-3 余弦跑 12000 步:看混合训练能不能学会。
|
| 12 |
+
|
| 13 |
+
结论(11-08,见 recall.md 的诊断一节):A、B 都能 100%;C 里只有 1e-3 余弦能学会 N=2
|
| 14 |
+
(第 8000 步 100%,loss 0.001),恒定 3e-3 跑 20000 步也停在 53%(loss ≈ ln2);D 12000 步仍 56%,
|
| 15 |
+
因为混合数据里 N=2 只占约 1/15,信号被稀释。
|
| 16 |
+
"""
|
| 17 |
+
import sys
|
| 18 |
+
import os
|
| 19 |
+
import time
|
| 20 |
+
|
| 21 |
+
import numpy as np
|
| 22 |
+
|
| 23 |
+
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
|
| 24 |
+
|
| 25 |
+
import zerocg as zc # noqa: E402
|
| 26 |
+
from zerocg import functional as F # noqa: E402
|
| 27 |
+
from models import TransformerLM # noqa: E402
|
| 28 |
+
|
| 29 |
+
NK, NV = 16, 16
|
| 30 |
+
QUERY = NK + NV
|
| 31 |
+
CFG = dict(n_embd=64, n_head=4, n_layer=2)
|
| 32 |
+
MAX_N = 16
|
| 33 |
+
TEST_N = [2, 4, 8, 12, 16]
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def batch(rng, B, n):
|
| 37 |
+
keys = np.stack([rng.permutation(NK)[:n] for _ in range(B)])
|
| 38 |
+
vals = rng.integers(NK, NK + NV, size=(B, n))
|
| 39 |
+
seq = np.empty((B, 2 * n + 2), dtype=np.int64)
|
| 40 |
+
seq[:, 0:2 * n:2], seq[:, 1:2 * n:2] = keys, vals
|
| 41 |
+
q = rng.integers(0, n, B)
|
| 42 |
+
seq[:, -2], seq[:, -1] = QUERY, keys[np.arange(B), q]
|
| 43 |
+
return seq, vals[np.arange(B), q]
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def model():
|
| 47 |
+
return TransformerLM(NK + NV + 1, 2 * MAX_N + 2, **CFG)
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def acc(model, tests):
|
| 51 |
+
with zc.no_grad():
|
| 52 |
+
return {n: round(float((model(tests[n][0]).data[:, -1].argmax(-1) == tests[n][1]).mean()), 3)
|
| 53 |
+
for n in tests}
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def tests():
|
| 57 |
+
rng = np.random.default_rng(7000)
|
| 58 |
+
return {n: batch(rng, 512, n) for n in TEST_N}
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def run(label, iters, lr0, lr1, warmup, n_fixed=None, seed=0, eval_every=2000):
|
| 62 |
+
np.random.seed(seed)
|
| 63 |
+
rng = np.random.default_rng(seed)
|
| 64 |
+
m = model()
|
| 65 |
+
opt = zc.optim.AdamW(m.parameters(), lr=lr0, weight_decay=0.01)
|
| 66 |
+
t = tests()
|
| 67 |
+
t0 = time.time()
|
| 68 |
+
for it in range(iters):
|
| 69 |
+
frac = it / max(1, iters - 1)
|
| 70 |
+
opt.lr = lr0 * min(1, (it + 1) / max(1, warmup)) * (lr1 / lr0) ** frac
|
| 71 |
+
n = n_fixed if n_fixed else int(rng.integers(2, MAX_N + 1))
|
| 72 |
+
x, y = batch(rng, 32, n)
|
| 73 |
+
loss = F.cross_entropy(m(x)[:, -1], y)
|
| 74 |
+
opt.zero_grad()
|
| 75 |
+
loss.backward()
|
| 76 |
+
zc.optim.clip_grad_norm(m.parameters(), 1.0)
|
| 77 |
+
opt.step()
|
| 78 |
+
if (it + 1) % eval_every == 0:
|
| 79 |
+
a = acc(m, t)
|
| 80 |
+
print(f" {label} step {it+1:5d} loss {float(loss.data):.3f} "
|
| 81 |
+
+ " ".join(f"N={k}:{a[k]}" for k in TEST_N) + f" ({time.time()-t0:.0f}s)", flush=True)
|
| 82 |
+
return a
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def main():
|
| 86 |
+
print("A 固定 32 条 N=2 反复训练(记忆)", flush=True)
|
| 87 |
+
np.random.seed(0)
|
| 88 |
+
m = model()
|
| 89 |
+
rng = np.random.default_rng(0)
|
| 90 |
+
x, y = batch(rng, 32, 2)
|
| 91 |
+
opt = zc.optim.AdamW(m.parameters(), lr=1e-3, weight_decay=0.01)
|
| 92 |
+
for it in range(600):
|
| 93 |
+
opt.lr = 3e-3 * min(1, (it + 1) / 100)
|
| 94 |
+
loss = F.cross_entropy(m(x)[:, -1], y)
|
| 95 |
+
opt.zero_grad()
|
| 96 |
+
loss.backward()
|
| 97 |
+
opt.step()
|
| 98 |
+
if (it + 1) % 200 == 0:
|
| 99 |
+
with zc.no_grad():
|
| 100 |
+
a = float((m(x).data[:, -1].argmax(-1) == y).mean())
|
| 101 |
+
print(f" mem step {it+1} loss {float(loss.data):.4f} acc {a}", flush=True)
|
| 102 |
+
|
| 103 |
+
print("B 只训 N=1(复制)", flush=True)
|
| 104 |
+
run("N1", 1500, 3e-3, 3e-3, 100, n_fixed=1, eval_every=500)
|
| 105 |
+
|
| 106 |
+
print("C1 只训 N=2,恒定 3e-3(recall.py 的原设置)", flush=True)
|
| 107 |
+
run("C1", 8000, 3e-3, 3e-3, 100, n_fixed=2)
|
| 108 |
+
print("C2 只训 N=2,余弦 3e-3→3e-4", flush=True)
|
| 109 |
+
run("C2", 8000, 3e-3, 3e-4, 200, n_fixed=2)
|
| 110 |
+
print("C3 只训 N=2,余弦 1e-3→1e-4", flush=True)
|
| 111 |
+
run("C3", 8000, 1e-3, 1e-4, 200, n_fixed=2)
|
| 112 |
+
|
| 113 |
+
print("D 混合 N,余弦 1e-3→1e-4,12000 步", flush=True)
|
| 114 |
+
run("D", 12000, 1e-3, 1e-4, 200, n_fixed=None)
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
if __name__ == "__main__":
|
| 118 |
+
main()
|
experiments/recall_prereg.md
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 预注册实验:精确回忆(联想检索),训练加长到 2 万步
|
| 2 |
+
|
| 3 |
+
> 写于实验之前(2026-11-01)。跑完在 recall.md 里看结果,这里不改。
|
| 4 |
+
|
| 5 |
+
**背景**:10-11 的 recall 实验只训 3000 步,Transformer 和 ZeroPoint 都接近瞎猜(N=16 时 14~18%,瞎猜 6%),无结论。
|
| 6 |
+
|
| 7 |
+
**做法**:同一脚本(recall.py),RECALL_ITERS=20000,其余不变(3 个模型 × 2 个种子,宽 64 / 2 层 / 4 头;N 在 2~16 随机;测 N = 2, 4, 8, 12, 16)。
|
| 8 |
+
|
| 9 |
+
**预测**
|
| 10 |
+
- P1 Transformer 学会:N=16 准确率 ≥ 80%(来源:注意力做键值检索是已知能力,3000 步不够是训练量问题)。
|
| 11 |
+
- P2 ZeroPoint 落后:N=16 时比 Transformer 低 ≥ 20 个百分点(来源:16 个槽、每头 16 维,要原样存 16 对键值超出容量;无依据的具体数字,仅作记录)。
|
| 12 |
+
- P3 槽数有用:32 槽比 16 槽在 N=16 时高 ≥ 5 个百分点(无依据,仅作记录)。
|
| 13 |
+
- 判定:2 个种子的均值;差别小于两个种子之间的差距时不算数。
|
experiments/results/20261011-1223-compare.md
ADDED
|
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
> **历史存档**(2026-10-11 生成,原名 ZeroMem)。本文件保留生成时原样,未按 2026-11-08 的复测更正;最新版本与更正见 [`../compare.md`](../compare.md) 和 [`../diagnose.md`](../diagnose.md)。
|
| 2 |
+
|
| 3 |
+
# ZeroMem vs Transformer 对照实验
|
| 4 |
+
|
| 5 |
+
- 语料:Python 标准库源码,1,558,395 字符,词表 96;90% 训练 / 10% 验证
|
| 6 |
+
- 三个模型:宽 64、2 层、4 头、Pre-LN、GELU MLP 4 倍;ZeroMem 每头 16 个槽、可学习初始状态、分块并行(块长 32)、无位置编码
|
| 7 |
+
- 训练:上下文 64、batch 32、2000 步、AdamW(权重衰减 0.01)、预热 100 步 + 余弦、梯度裁剪 1.0
|
| 8 |
+
- 学习率:每个模型在 [0.001, 0.003, 0.006] 里各试一次(种子 0),取最好的再跑种子 [1, 2]
|
| 9 |
+
- 验证 loss = 40 个固定验证 batch 的平均交叉熵(nats/字符,越低越好),均值 ± 标准差(3 个种子)
|
| 10 |
+
- 计算量 = 前向时每个字的矩阵乘浮点运算次数(不受机器负载影响;训练一步约 3 倍)
|
| 11 |
+
|
| 12 |
+
| 模型 | 参数量 | 最佳学习率 | 验证 loss | 前向计算量(每字) | 推理状态(第 64 字 / 第 4096 字,浮点数) |
|
| 13 |
+
|---|---|---|---|---|---|
|
| 14 |
+
| Transformer | 116,576 | 0.006 | 1.5452 ± 0.0100 | 242 K | 16,384 / 1,048,576(KV 缓存) |
|
| 15 |
+
| Transformer (sin) | 112,480 | 0.006 | 1.5271 ± 0.0083 | 242 K | 16,384 / 1,048,576(KV 缓存) |
|
| 16 |
+
| ZeroMem | 115,240 | 0.006 | 1.5214 ± 0.0134 | 232 K | 2,048 / 2,048(固定) |
|
| 17 |
+
|
| 18 |
+
## 学习率搜索(种子 0 的验证 loss)
|
| 19 |
+
|
| 20 |
+
| 模型 | lr=0.001 | lr=0.003 | lr=0.006 |
|
| 21 |
+
|---|---|---|---|
|
| 22 |
+
| Transformer | 1.6908 | 1.5530 | 1.5520 |
|
| 23 |
+
| Transformer (sin) | 1.8006 | 1.5939 | 1.5374 |
|
| 24 |
+
| ZeroMem | 1.6682 | 1.5433 | 1.5311 |
|
| 25 |
+
|
| 26 |
+
## 长度外推(训练上下文 64;8 条验证序列 × 3 个种子平均)
|
| 27 |
+
|
| 28 |
+
| 位置 | ZeroMem | Transformer (sin) | Transformer(可学习位置编码) |
|
| 29 |
+
|---|---|---|---|
|
| 30 |
+
| 0 ~ 63(训练见过的长度) | 1.5550 | 1.4512 | 能跑 |
|
| 31 |
+
| 64 ~ 255 | 1.6164 | 4.6359 | 跑不了(只有 64 个位置) |
|
| 32 |
+
| 256 ~ 1023 | 1.6410 | 4.6057 | 跑不了(只有 64 个位置) |
|
| 33 |
+
| 1024 ~ 4095 | 1.7040 | —(注意力矩阵太大,没测) | 跑不了(只有 64 个位置) |
|
| 34 |
+
|
| 35 |
+
## ZeroMem 槽的使用情况(有效槽数,满分 16)
|
| 36 |
+
|
| 37 |
+
有效槽数 = 平均读/写地址分布的 exp(熵):16 个槽用得一样多 = 16,全挤在一个槽 = 1。
|
| 38 |
+
|
| 39 |
+
| 层 | 读 | 写 |
|
| 40 |
+
|---|---|---|
|
| 41 |
+
| 0 | 10.1 | 3.2 |
|
| 42 |
+
| 1 | 9.9 | 5.0 |
|
experiments/results/20261011-1258-ablation.md
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
> **历史存档**(2026-10-11 生成,原名 ZeroMem)。本文件保留生成时原样,未按 2026-11-08 的复测更正;最新版本见 [`../ablation.md`](../ablation.md)。
|
| 2 |
+
|
| 3 |
+
# ZeroMem 消融实验(仿论文表 3)
|
| 4 |
+
|
| 5 |
+
每行只改一处,其余和 base 相同。2000 步,lr 0.003,种子 [0, 1, 2],其余训练设置同 [`compare.md`](../compare.md)。
|
| 6 |
+
参数量和 base 差超过 1% 的变体,自动调整 MLP 宽度对齐参数量。
|
| 7 |
+
Δ = 该行验证 loss 均值 − base 均值(**正数表示变差**)。
|
| 8 |
+
|
| 9 |
+
> 判断标准:base 的种子间标准差是 0.0073。|Δ| 小于约 2 倍标准差(0.0145)的差别**不能算数**。
|
| 10 |
+
|
| 11 |
+
| | 改动 | 参数量 | 验证 loss | Δ | 结论 |
|
| 12 |
+
|---|---|---|---|---|---|
|
| 13 |
+
| base | ZeroMem:4 头 × 16 槽,写强度门,写入竞争,分层初始化,λ 可学,可学初始状态,无位置编码 | 115,240 | 1.5500 ± 0.0073 | — | — |
|
| 14 |
+
| (A) | 1 个头(每头 64 维)(MLP 宽度调为 306 以对齐参数) | 115,366 | 1.5776 ± 0.0036 | +0.0276 | 变差:这个部件有用 |
|
| 15 |
+
| (A) | 8 个头(每头 8 维)(MLP 宽度调为 189 以对齐参数) | 115,210 | 1.5673 ± 0.0067 | +0.0173 | 变差:这个部件有用 |
|
| 16 |
+
| (B) | 每头 4 个槽(MLP 宽度调为 311 以对齐参数) | 115,318 | 1.5466 ± 0.0087 | -0.0034 | 差别不显著 |
|
| 17 |
+
| (B) | 每头 32 个槽(MLP 宽度调为 183 以对齐参数) | 115,222 | 1.5624 ± 0.0055 | +0.0124 | 差别不显著 |
|
| 18 |
+
| (C) | 去掉写强度门 | 114,720 | 1.5522 ± 0.0072 | +0.0022 | 差别不显著 |
|
| 19 |
+
| (D) | 写入不竞争:softmax 换成每槽独立的 sigmoid | 115,240 | 1.5487 ± 0.0066 | -0.0013 | 差别不显著 |
|
| 20 |
+
| (E) | 不分层:所有槽初始半衰期相同(约 9.8 个字) | 115,240 | 1.6010 ± 0.0079 | +0.0510 | 变差:这个部件有用 |
|
| 21 |
+
| (F) | λ 固定在初始值,不学习 | 115,112 | 1.5784 ± 0.0097 | +0.0284 | 变差:这个部件有用 |
|
| 22 |
+
| (G) | 加上论文 §3.5 的正弦位置编码 | 115,240 | 1.6106 ± 0.0078 | +0.0606 | 变差:这个部件有用 |
|
experiments/results/20261011-1305-recall.md
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
> **历史存档**(3,000 步版本)。结论见 [`../recall.md`](../recall.md):学习率配置导致**无结论**。
|
| 2 |
+
|
| 3 |
+
# 长距离精确回忆(联想检索)
|
| 4 |
+
|
| 5 |
+
`k1 v1 k2 v2 … kN vN [问] kq` → 回答 kq 对应的值。16 个键、16 个值,随机猜的准确率是 1/16 ≈ 6%。
|
| 6 |
+
三个模型宽 64、2 层、4 头;训练 3000 步,N 在 2~16 随机;每格是 512 道测试题 × 种子 [0, 1] 的平均准确率。
|
| 7 |
+
|
| 8 |
+
| 模型 | N = 2 | N = 4 | N = 8 | N = 12 | N = 16 |
|
| 9 |
+
|---|---|---|---|---|---|
|
| 10 |
+
| Transformer | 52% | 32% | 23% | 18% | 14% |
|
| 11 |
+
| ZeroPoint (16 槽) | 51% | 32% | 24% | 21% | 17% |
|
| 12 |
+
| ZeroPoint (32 槽) | 51% | 31% | 25% | 20% | 18% |
|
experiments/results/20261018-0838-ema_longrun.md
ADDED
|
@@ -0,0 +1,57 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
> **历史存档**。其中"compare.py 里 ZeroPoint 用的是递推算法、和 Transformer 不对等"的说法有误:分块并行与递推形式数学等价(float32 差 1.4e-5),10-18 的修正实际是把 8 条验证序列增加到 32 条。
|
| 2 |
+
|
| 3 |
+
# 预注册实验:长序列变差与 ema 写入
|
| 4 |
+
|
| 5 |
+
> 写于实验之前(2026-10-18)。跑完后只在下面追加结果,不改这一节。
|
| 6 |
+
|
| 7 |
+
**观察**:compare.md 里 ZeroPoint(无 ema)在 1024~4095 字处的 loss(1.70)高于 64~255 字处(1.62)。
|
| 8 |
+
|
| 9 |
+
**假设 H1**:原因是长期槽(λ 接近 1)的数值随序列长度无界累积,超出训练时(64 字)见过的范围。
|
| 10 |
+
|
| 11 |
+
**预测**:
|
| 12 |
+
- ① 直接测:推理时槽值的最大绝对值,无 ema 的模型在 4000 字处明显大于 64 字处(比值 > 2);有 ema 的模型比值接近 1(< 1.5)。
|
| 13 |
+
- ② 间接测:开 ema 重新训练(其余完全相同),1024~4095 字处的 loss 不再明显高于 64~255 字处:差值 < 2 × 种子间标准差。
|
| 14 |
+
|
| 15 |
+
**H1 不成立时会看到**:① 无 ema 的槽值没有明显增长;或 ② 开 ema 后远处 loss 照样升高。
|
| 16 |
+
|
| 17 |
+
**判定**:① 和 ② 都成立 → 接受 H1,且 ema 是有效的修法;只有 ② 成立 → ema 有效但原因另有其他;
|
| 18 |
+
都不成立 → 放弃 H1,下一步查学到的 λ 值。另外要求 ema 不损害短距离:64 字处验证 loss 的差别 < 2 × 标准差。
|
| 19 |
+
|
| 20 |
+
**对照**:两组用同样的 3 个种子、同样 2000 步、学习率 0.006、同样语料;外推用 32 条 4096 字的验证序列,
|
| 21 |
+
两组都用分块并行算法算 loss(compare.py 里 ZeroPoint 用的是递推算法,和 Transformer 不对等,这里统一)。
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
## 结果(实验之后写)
|
| 25 |
+
|
| 26 |
+
### 64 字处的验证 loss(均值 ± 标准差,3 个种子)
|
| 27 |
+
|
| 28 |
+
| 模型 | 验证 loss |
|
| 29 |
+
|---|---|
|
| 30 |
+
| no-ema | 1.5214 ± 0.0134 |
|
| 31 |
+
| ema | 1.5182 ± 0.0076 |
|
| 32 |
+
|
| 33 |
+
### 长度外推(32 条 × 4096 字,分块并行算法;均值 ± 标准差)
|
| 34 |
+
|
| 35 |
+
| 位置 | no-ema | ema |
|
| 36 |
+
|---|---|---|
|
| 37 |
+
| 0 ~ 63 | 1.4372 ± 0.0187 | 1.4308 ± 0.0039 |
|
| 38 |
+
| 64 ~ 255 | 1.4742 ± 0.0341 | 1.6476 ± 0.0529 |
|
| 39 |
+
| 256 ~ 1023 | 1.6617 ± 0.0234 | 3.1084 ± 0.2157 |
|
| 40 |
+
| 1024 ~ 4095 | 1.6932 ± 0.0305 | 3.4153 ± 0.2402 |
|
| 41 |
+
|
| 42 |
+
### 槽值的最大绝对值(递推推理,3 个种子平均)
|
| 43 |
+
|
| 44 |
+
| 位置 | no-ema | ema |
|
| 45 |
+
|---|---|---|
|
| 46 |
+
| 64 | 9.08 | 1.60 |
|
| 47 |
+
| 256 | 10.07 | 2.34 |
|
| 48 |
+
| 1024 | 9.88 | 1.51 |
|
| 49 |
+
| 4000 | 16.07 | 2.27 |
|
| 50 |
+
|
| 51 |
+
## 判定(按预注册的标准)
|
| 52 |
+
|
| 53 |
+
- 预测 ①(槽值比值 4000 字 / 64 字):无 ema 1.93,有 ema 1.44 → **不成立**
|
| 54 |
+
- 预测 ②(ema 模型 1024~4095 与 64~255 的 loss 差):+1.7677,噪声阈值 0.4805 → **不成立**(对照:无 ema 的差是 +0.2190)
|
| 55 |
+
- ema 不损害短距离:64 字处 loss 差 -0.0033 → **成立**
|
| 56 |
+
|
| 57 |
+
**结论**:① ② 都不成立:**放弃 H1**,下一步查学到的 λ 值
|
experiments/results/20261018-0901-diagnose.md
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
> **历史存档**。其中"有效槽数 3~5"是未归一化门控分布的偏低值;归一化后读约 10/16、写约 6~10/层(见 [`../diagnose.md`](../diagnose.md))。
|
| 2 |
+
|
| 3 |
+
# 预注册实验:查三个问题
|
| 4 |
+
|
| 5 |
+
> 写于实验之前(2026-10-18)。跑完只在下面追加结果,不改这一节。
|
| 6 |
+
|
| 7 |
+
**背景**:ema_longrun.md 否定了"槽值无界累积"假设。无 ema 的模型在 1024~4095 字处 loss 比 64~255 字处高约 0.22;
|
| 8 |
+
有 ema 的高约 1.77。compare.md 里 16 个槽只有 3~5 个被写入。
|
| 9 |
+
|
| 10 |
+
## 假设和预测
|
| 11 |
+
|
| 12 |
+
**H2(训练长度)**:长序列变差是因为训练只见过 64 字;模型在更长的位置上遇到的槽状态分布是训练时没见过的。
|
| 13 |
+
- 预测:训练长度改为 256(每步字数不变、其余相同)后,1024~4095 与 64~255 的 loss 差 **缩小到 T64 的一半以下**。
|
| 14 |
+
- 不成立时会看到:差距没有缩小(< 2 × 标准差的变化)。
|
| 15 |
+
|
| 16 |
+
**H3(ema 为什么更差)**:ema 把长期槽的写入乘上 (1-λ),训练时(64 字)长期槽几乎不起作用,模型学会只用短期槽;
|
| 17 |
+
到了长序列,读地址仍指向短期槽,但短期槽里的内容和训练时的统计不同。
|
| 18 |
+
- 预测 ③a:ema 模型的**写地址**集中在半衰期短的槽上:半衰期 > 16 字的槽,写入占比 < T64 模型的一半。
|
| 19 |
+
- 预测 ③b:ema 模型学到的长期槽 λ 更接近 1(半衰期更长),或者长期槽的读占比在远处位置上升。
|
| 20 |
+
- 不成立时会看到:两组的写入分布没有明显差别。
|
| 21 |
+
|
| 22 |
+
**H4(槽塌缩)**:塌缩不是训练失败,而是写地址的 softmax 竞争让少数槽赢者通吃;在不同序列位置上"赢的槽"不同,
|
| 23 |
+
所以整体统计(compare.md 的 3~5 个有效槽)低估了实际用到的槽数。
|
| 24 |
+
- 预测 ④:按位置分段(0~63 / 64~255 / 256~1023 / 1024~4095)分别统计有效写槽数,各段的并集 > 整体有效槽数 × 1.5;
|
| 25 |
+
同时各段的读/写分布之间有明显差别(任意两段的 JS 散度 > 0.1)。
|
| 26 |
+
- 不成立时会看到:各段的分布几乎相同,用的始终是同几个槽。
|
| 27 |
+
|
| 28 |
+
**判定**:每条按上面的数值阈值判,3 个种子均值;差别和 2 × 种子间标准差比。
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
## 结果(实验之后写)
|
| 32 |
+
|
| 33 |
+
### 验证 loss(64 字)和长度外推(32 条 × 4096 字;均值 ± 标准差,3 个种子)
|
| 34 |
+
|
| 35 |
+
| 条件 | 64 字验证 | 0~63 | 64~255 | 256~1023 | 1024~4095 | 远−近 差 |
|
| 36 |
+
|---|---|---|---|---|---|---|
|
| 37 |
+
| T64 | 1.5214 ± 0.0134 | 1.4372 ± 0.0187 | 1.4742 ± 0.0341 | 1.6617 ± 0.0234 | 1.6932 ± 0.0305 | +0.2190 ± 0.0630 |
|
| 38 |
+
| T64-ema | 1.5182 ± 0.0076 | 1.4308 ± 0.0039 | 1.6476 ± 0.0529 | 3.1084 ± 0.2157 | 3.4153 ± 0.2402 | +1.7677 ± 0.2188 |
|
| 39 |
+
| T256 | 1.5572 ± 0.0060 | 1.4741 ± 0.0117 | 1.4486 ± 0.0098 | 1.5098 ± 0.0022 | 1.5155 ± 0.0048 | +0.0669 ± 0.0051 |
|
| 40 |
+
|
| 41 |
+
### H3:写入落在半衰期 > 16 字的槽上的占比(0~63 段),和学到的最长半衰期
|
| 42 |
+
|
| 43 |
+
| 条件 | 长期槽写入占比 | 最长半衰期(字) | 远处读长期槽占比的变化 |
|
| 44 |
+
|---|---|---|---|
|
| 45 |
+
| T64 | 0.061 ± 0.005 | 109 | +0.009 |
|
| 46 |
+
| T64-ema | 0.014 ± 0.001 | 103 | +0.007 |
|
| 47 |
+
|
| 48 |
+
### H4:槽的使用(T64,写地址;各层各头平均)
|
| 49 |
+
|
| 50 |
+
| 段 | 有效写槽数 |
|
| 51 |
+
|---|---|
|
| 52 |
+
| 0~63 | 8.6 |
|
| 53 |
+
| 64~255 | 8.7 |
|
| 54 |
+
| 256~1023 | 8.5 |
|
| 55 |
+
| 1024~4095 | 8.9 |
|
| 56 |
+
| 整体 | 8.7 |
|
| 57 |
+
|
| 58 |
+
各段用到的槽的并集 ÷ 整体有效槽数 = 0.68;各段分布之间最大的 JS 散度 = 0.010
|
| 59 |
+
|
| 60 |
+
## 判定(按预注册的标准)
|
| 61 |
+
|
| 62 |
+
- H2 训练长度:远−近差 T64 +0.2190 → T256 +0.0669 → **成立**
|
| 63 |
+
- H3a ema 写入避开长期槽:占比 0.061 → 0.014 → **成立**
|
| 64 |
+
- H3b ema 的长期槽更长或远处更依赖长期槽 → **不成立**
|
| 65 |
+
- H4 槽塌缩是分段使用不同的槽:并集比 0.68,JS 0.010 → **不成立**
|
| 66 |
+
|
| 67 |
+
模型已保存在 experiments/out/diag_*.npz,可以事后再分析。
|
experiments/results/20261101-0814-slots.md
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
> **历史存档**(与 [`../slots.md`](../slots.md) 内容相同)。
|
| 2 |
+
|
| 3 |
+
# 预注册分析:长期槽在干什么
|
| 4 |
+
|
| 5 |
+
> 写于分析之前(2026-11-01)。跑完只追加结果,不改这一节。用 experiments/out/diag_*.npz(T64 / T64-ema / T256 各 3 种子)。
|
| 6 |
+
|
| 7 |
+
**矛盾**:T64 模型写入只有 6.1% 落在半衰期 > 16 字的槽;但消融里去掉分层初始化变差 0.051。
|
| 8 |
+
|
| 9 |
+
**三个可区分的假设**
|
| 10 |
+
- H9 读取端在用:长期槽虽然写得少,但**读地址**落在长期槽上的占比明显高于写占比(读占比 ≥ 2 × 写占比)。
|
| 11 |
+
→ 长期槽像"慢变的背景",少量写入、持续读取。
|
| 12 |
+
- H10 中期槽才是关键:分层初始化的作用主要来自半衰期 4~64 字的槽,不是 > 64 的;读写都集中在 4~64 桶(该桶读 + 写占比合计 > 50%)。
|
| 13 |
+
- H11 训练长度决定长期槽的使用:T256 模型在 > 16 字桶上的写占比 ≥ 2 × T64 的。
|
| 14 |
+
来源:diagnose 里 T256 远处 loss 明显更好,推测它学会了用更长的槽。
|
| 15 |
+
|
| 16 |
+
**判定**:3 个种子平均;各桶占比按所有层、所有头的平均。H9 / H11 的"2 倍"是无依据的整数阈值,仅作记录;
|
| 17 |
+
同时报告原始占比,让读者自己看。
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
## 结果(分析之后写)
|
| 21 |
+
|
| 22 |
+
验证集 4 条 × 2048 字递推;占比 = 该桶内所有槽的地址权重之和(读 / 写各自归一化),层和头平均,3 个种子平均。
|
| 23 |
+
|
| 24 |
+
| 条件 | 量 | 半衰期 0~4 | 半衰期 4~16 | 半衰期 16~64 | 半衰期 64~∞ |
|
| 25 |
+
|---|---|---|---|---|---|
|
| 26 |
+
| T64 | 读 | 0.703 | 0.120 | 0.127 | 0.050 |
|
| 27 |
+
| T64 | 写 | 0.815 | 0.122 | 0.058 | 0.006 |
|
| 28 |
+
| T64-ema | 读 | 0.524 | 0.149 | 0.167 | 0.159 |
|
| 29 |
+
| T64-ema | 写 | 0.925 | 0.060 | 0.014 | 0.001 |
|
| 30 |
+
| T256 | 读 | 0.660 | 0.180 | 0.152 | 0.008 |
|
| 31 |
+
| T256 | 写 | 0.696 | 0.167 | 0.135 | 0.001 |
|
| 32 |
+
|
| 33 |
+
## 判定(按预注册的标准)
|
| 34 |
+
|
| 35 |
+
- H9 读取端在用长期槽:> 16 字桶 读 0.177 vs 写 0.064 → **成立**(2 倍阈值无依据,仅作记录)
|
| 36 |
+
- H10 中期槽是关键:4~64 字桶 读写平均占比 0.213 > 0.5 → **不成立**
|
| 37 |
+
- H11 训练长度 256 更用长期槽:> 16 字桶写占比 T256 0.136 vs T64 0.064 → **成立**(2 倍阈值无依据,仅作记录)
|
experiments/results/20261101-0936-recall.md
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
> **历史存档**(20,000 步版本)。结论见 [`../recall.md`](../recall.md):学习率配置导致**无结论**。
|
| 2 |
+
|
| 3 |
+
# 长距离精确回忆(联想检索)
|
| 4 |
+
|
| 5 |
+
`k1 v1 k2 v2 … kN vN [问] kq` → 回答 kq 对应的值。16 个键、16 个值,随机猜的准确率是 1/16 ≈ 6%。
|
| 6 |
+
三个模型宽 64、2 层、4 头;训练 20000 步,N 在 2~16 随机;每格是 512 道测试题 × 种子 [0, 1] 的平均准确率。
|
| 7 |
+
|
| 8 |
+
| 模型 | N = 2 | N = 4 | N = 8 | N = 12 | N = 16 |
|
| 9 |
+
|---|---|---|---|---|---|
|
| 10 |
+
| Transformer | 50% | 30% | 23% | 20% | 18% |
|
| 11 |
+
| ZeroPoint (16 槽) | 51% | 32% | 23% | 20% | 16% |
|
| 12 |
+
| ZeroPoint (32 槽) | 49% | 31% | 24% | 20% | 17% |
|
experiments/slots.md
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 预注册分析:长期槽在干什么
|
| 2 |
+
|
| 3 |
+
> 写于分析之前(2026-11-01)。跑完只追加结果,不改这一节。用 experiments/out/diag_*.npz(T64 / T64-ema / T256 各 3 种子)。
|
| 4 |
+
|
| 5 |
+
**矛盾**:T64 模型写入只有 6.1% 落在半衰期 > 16 字的槽;但消融里去掉分层初始化变差 0.051。
|
| 6 |
+
|
| 7 |
+
**三个可区分的假设**
|
| 8 |
+
- H9 读取端在用:长期槽虽然写得少,但**读地址**落在长期槽上的占比明显高于写占比(读占比 ≥ 2 × 写占比)。
|
| 9 |
+
→ 长期槽像"慢变的背景",少量写入、持续读取。
|
| 10 |
+
- H10 中期槽才是关键:分层初始化的作用主要来自半衰期 4~64 字的槽,不是 > 64 的;读写都集中在 4~64 桶(该桶读 + 写占比合计 > 50%)。
|
| 11 |
+
- H11 训练长度决定长期槽的使用:T256 模型在 > 16 字桶上的写占比 ≥ 2 × T64 的。
|
| 12 |
+
来源:diagnose 里 T256 远处 loss 明显更好,推测它学会了用更长的槽。
|
| 13 |
+
|
| 14 |
+
**判定**:3 个种子平均;各桶占比按所有层、所有头的平均。H9 / H11 的"2 倍"是无依据的整数阈值,仅作记录;
|
| 15 |
+
同时报告原始占比,让读者自己看。
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
## 结果(分析之后写)
|
| 19 |
+
|
| 20 |
+
验证集 4 条 × 2048 字递推;占比 = 该桶内所有槽的地址权重之和(读 / 写各自归一化),层和头平均,3 个种子平均。
|
| 21 |
+
|
| 22 |
+
| 条件 | 量 | 半衰期 0~4 | 半衰期 4~16 | 半衰期 16~64 | 半衰期 64~∞ |
|
| 23 |
+
|---|---|---|---|---|---|
|
| 24 |
+
| T64 | 读 | 0.703 | 0.120 | 0.127 | 0.050 |
|
| 25 |
+
| T64 | 写 | 0.815 | 0.122 | 0.058 | 0.006 |
|
| 26 |
+
| T64-ema | 读 | 0.524 | 0.149 | 0.167 | 0.159 |
|
| 27 |
+
| T64-ema | 写 | 0.925 | 0.060 | 0.014 | 0.001 |
|
| 28 |
+
| T256 | 读 | 0.660 | 0.180 | 0.152 | 0.008 |
|
| 29 |
+
| T256 | 写 | 0.696 | 0.167 | 0.135 | 0.001 |
|
| 30 |
+
|
| 31 |
+
## 判定(按预注册的标准)
|
| 32 |
+
|
| 33 |
+
- H9 读取端在用长期槽:> 16 字桶 读 0.177 vs 写 0.064 → **成立**(2 倍阈值无依据,仅作记录)
|
| 34 |
+
- H10 中期槽是关键:4~64 字桶 读写平均占比 0.213 > 0.5 → **不成立**
|
| 35 |
+
- H11 训练长度 256 更用长期槽:> 16 字桶写占比 T256 0.136 vs T64 0.064 → **成立**(2 倍阈值无依据,仅作记录)
|
experiments/slots.py
ADDED
|
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""预注册分析:长期槽到底在干什么?(用 diagnose 保存的 9 个模型,不重训)
|
| 2 |
+
|
| 3 |
+
python3 experiments/slots.py
|
| 4 |
+
|
| 5 |
+
背景:diagnose 发现无 ema 的模型只有 6% 的写入落在半衰期 > 16 字的槽上,但消融说明"分层初始化"有用(去掉变差 0.051)。
|
| 6 |
+
结果:experiments/slots.md + results/ 存档。
|
| 7 |
+
"""
|
| 8 |
+
import os
|
| 9 |
+
|
| 10 |
+
import numpy as np
|
| 11 |
+
|
| 12 |
+
from common import HERE, batches, prepare_data, save_result
|
| 13 |
+
from models import ZeroX
|
| 14 |
+
|
| 15 |
+
OUT = os.path.join(HERE, "out")
|
| 16 |
+
SEEDS = [0, 1, 2]
|
| 17 |
+
T_RUN = 2048
|
| 18 |
+
BUCKETS = [(0, 4), (4, 16), (16, 64), (64, 1e9)] # 半衰期分桶(字)
|
| 19 |
+
|
| 20 |
+
PREREG = """# 预注册分析:长期槽在干什么
|
| 21 |
+
|
| 22 |
+
> 写于分析之前(2026-11-01)。跑完只追加结果,不改这一节。用 experiments/out/diag_*.npz(T64 / T64-ema / T256 各 3 种子)。
|
| 23 |
+
|
| 24 |
+
**矛盾**:T64 模型写入只有 6.1% 落在半衰期 > 16 字的槽;但消融里去掉分层初始化变差 0.051。
|
| 25 |
+
|
| 26 |
+
**三个可区分的假设**
|
| 27 |
+
- H9 读取端在用:长期槽虽然写得少,但**读地址**落在长期槽上的占比明显高于写占比(读占比 ≥ 2 × 写占比)。
|
| 28 |
+
→ 长期槽像"慢变的背景",少量写入、持续读取。
|
| 29 |
+
- H10 中期槽才是关键:分层初始化的作用主要来自半衰期 4~64 字的槽,不是 > 64 的;读写都集中在 4~64 桶(该桶读 + 写占比合计 > 50%)。
|
| 30 |
+
- H11 训练长度决定长期槽的使用:T256 模型在 > 16 字桶上的写占比 ≥ 2 × T64 的。
|
| 31 |
+
来源:diagnose 里 T256 远处 loss 明显更好,推测它学会了用更长的槽。
|
| 32 |
+
|
| 33 |
+
**判定**:3 个种子平均;各桶占比按所有层、所有头的平均。H9 / H11 的"2 倍"是无依据的整数阈值,仅作记录;
|
| 34 |
+
同时报告原始占比,让读者自己看。
|
| 35 |
+
"""
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def main():
|
| 39 |
+
text, V, train_data, val_data = prepare_data()
|
| 40 |
+
x, _ = next(batches(val_data, 1, T_RUN, 7, batch=4))
|
| 41 |
+
cfg = dict(n_embd=64, n_head=4, n_layer=2, n_slot=16)
|
| 42 |
+
rows = {}
|
| 43 |
+
for cond in ("T64", "T64-ema", "T256"):
|
| 44 |
+
reads, writes = np.zeros(len(BUCKETS)), np.zeros(len(BUCKETS))
|
| 45 |
+
for seed in SEEDS:
|
| 46 |
+
model = ZeroX(V, ema=(cond == "T64-ema"), **cfg)
|
| 47 |
+
model.load(os.path.join(OUT, f"diag_{cond}_s{seed}.npz")).eval()
|
| 48 |
+
states = model.init_state(4)
|
| 49 |
+
r_acc = [np.zeros((b.mix.H, b.mix.M)) for b in model.blocks]
|
| 50 |
+
w_acc = [np.zeros((b.mix.H, b.mix.M)) for b in model.blocks]
|
| 51 |
+
for t in range(T_RUN):
|
| 52 |
+
_, states = model.step(x[:, t], states, t)
|
| 53 |
+
for l, b in enumerate(model.blocks):
|
| 54 |
+
a, w = b.mix.last_addr
|
| 55 |
+
r_acc[l] += a.mean(0)
|
| 56 |
+
w_acc[l] += w.mean(0)
|
| 57 |
+
for l, b in enumerate(model.blocks):
|
| 58 |
+
hl = np.log(0.5) / np.log(b.mix.retention().data) # (H, M)
|
| 59 |
+
r = r_acc[l] / r_acc[l].sum(-1, keepdims=True)
|
| 60 |
+
w = w_acc[l] / w_acc[l].sum(-1, keepdims=True)
|
| 61 |
+
for k, (lo, hi) in enumerate(BUCKETS):
|
| 62 |
+
m = (hl >= lo) & (hl < hi)
|
| 63 |
+
reads[k] += (r * m).sum() / (r.shape[0]) / len(model.blocks)
|
| 64 |
+
writes[k] += (w * m).sum() / (w.shape[0]) / len(model.blocks)
|
| 65 |
+
rows[cond] = (reads / len(SEEDS), writes / len(SEEDS))
|
| 66 |
+
|
| 67 |
+
lines = [PREREG, "\n## 结果(分析之后写)\n",
|
| 68 |
+
f"验证集 4 条 × {T_RUN} 字递推;占比 = 该桶内所有槽的地址权重之和(读 / 写各自归一化),层和头平均,3 个种子平均。\n",
|
| 69 |
+
"| 条件 | 量 | " + " | ".join(f"半衰期 {lo}~{'∞' if hi > 1e8 else int(hi)}" for lo, hi in BUCKETS) + " |",
|
| 70 |
+
"|---|---|" + "---|" * len(BUCKETS)]
|
| 71 |
+
for cond, (r, w) in rows.items():
|
| 72 |
+
lines.append(f"| {cond} | 读 | " + " | ".join(f"{v:.3f}" for v in r) + " |")
|
| 73 |
+
lines.append(f"| {cond} | 写 | " + " | ".join(f"{v:.3f}" for v in w) + " |")
|
| 74 |
+
r64, w64 = rows["T64"]
|
| 75 |
+
r256, w256 = rows["T256"]
|
| 76 |
+
long_r, long_w = r64[2] + r64[3], w64[2] + w64[3]
|
| 77 |
+
h9 = long_r >= 2 * long_w
|
| 78 |
+
h10 = (r64[1] + r64[2] + w64[1] + w64[2]) / 2 > 0.5
|
| 79 |
+
h11 = (w256[2] + w256[3]) >= 2 * long_w
|
| 80 |
+
lines += ["", "## 判定(按预注册的标准)\n",
|
| 81 |
+
f"- H9 读取端在用长期槽:> 16 字桶 读 {long_r:.3f} vs 写 {long_w:.3f} → **{'成立' if h9 else '不成立'}**(2 倍阈值无依据,仅作记录)",
|
| 82 |
+
f"- H10 中期槽是关键:4~64 字桶 读写平均占比 {(r64[1] + r64[2] + w64[1] + w64[2]) / 2:.3f} > 0.5 → **{'成立' if h10 else '不成立'}**",
|
| 83 |
+
f"- H11 训练长度 256 更用长期槽:> 16 字桶写占比 T256 {w256[2] + w256[3]:.3f} vs T64 {long_w:.3f} → **{'成立' if h11 else '不成立'}**(2 倍阈值无依据,仅作记录)"]
|
| 84 |
+
save_result("slots", lines)
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
if __name__ == "__main__":
|
| 88 |
+
main()
|
models/__init__.py
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from .transformer import TransformerLM
|
| 2 |
+
from .zeropoint import MemoryMixer, ZeroX
|
| 3 |
+
|
| 4 |
+
__all__ = ["ZeroX", "MemoryMixer", "TransformerLM"]
|
models/transformer.py
ADDED
|
@@ -0,0 +1,63 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""对照组:标准的 GPT 式 Transformer(因果自注意力),结构上和 ZeroX 一一对应。
|
| 2 |
+
|
| 3 |
+
位置编码两种可选(《Attention Is All You Need》§3.5 两种都试过,结果几乎一样):
|
| 4 |
+
pos="learned" 可学习的位置向量,只有 block_size 个,处理不了更长的序列
|
| 5 |
+
pos="sin" 论文用的正弦编码,任意长度都能算(论文选它就是因为"可能可以外推")
|
| 6 |
+
"""
|
| 7 |
+
import numpy as np
|
| 8 |
+
|
| 9 |
+
from zerocg import Tensor
|
| 10 |
+
from zerocg import functional as F
|
| 11 |
+
from zerocg import nn
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def sinusoidal_encoding(T, d):
|
| 15 |
+
"""PE(pos, 2i) = sin(pos / 10000^(2i/d)),PE(pos, 2i+1) = cos(...)。论文 §3.5。"""
|
| 16 |
+
pos = np.arange(T)[:, None]
|
| 17 |
+
div = 10000 ** (np.arange(0, d, 2) / d)
|
| 18 |
+
pe = np.zeros((T, d), dtype=np.float32)
|
| 19 |
+
pe[:, 0::2] = np.sin(pos / div)
|
| 20 |
+
pe[:, 1::2] = np.cos(pos / div)
|
| 21 |
+
return pe
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class Block(nn.Module):
|
| 25 |
+
def __init__(self, n_embd, n_head, dropout=0.0):
|
| 26 |
+
self.ln1 = nn.LayerNorm(n_embd)
|
| 27 |
+
self.attn = nn.CausalSelfAttention(n_embd, n_head, dropout)
|
| 28 |
+
self.ln2 = nn.LayerNorm(n_embd)
|
| 29 |
+
self.fc1 = nn.Linear(n_embd, 4 * n_embd)
|
| 30 |
+
self.fc2 = nn.Linear(4 * n_embd, n_embd)
|
| 31 |
+
self.drop = nn.Dropout(dropout)
|
| 32 |
+
|
| 33 |
+
def forward(self, x):
|
| 34 |
+
x = x + self.attn(self.ln1(x))
|
| 35 |
+
return x + self.drop(self.fc2(F.gelu(self.fc1(self.ln2(x)))))
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
class TransformerLM(nn.Module):
|
| 39 |
+
def __init__(self, vocab_size, block_size, n_embd=64, n_head=4, n_layer=2, dropout=0.0, pos="learned"):
|
| 40 |
+
assert pos in ("learned", "sin")
|
| 41 |
+
self.block_size = block_size
|
| 42 |
+
self.pos_kind = pos
|
| 43 |
+
self.tok = nn.Embedding(vocab_size, n_embd)
|
| 44 |
+
if pos == "learned":
|
| 45 |
+
self.pos = nn.Embedding(block_size, n_embd)
|
| 46 |
+
self.blocks = [Block(n_embd, n_head, dropout) for _ in range(n_layer)]
|
| 47 |
+
self.ln_f = nn.LayerNorm(n_embd)
|
| 48 |
+
self.head = nn.Linear(n_embd, vocab_size)
|
| 49 |
+
|
| 50 |
+
def forward(self, idx):
|
| 51 |
+
T = idx.shape[1]
|
| 52 |
+
if self.pos_kind == "learned":
|
| 53 |
+
assert T <= self.block_size, f"位置编码只有 {self.block_size} 个,处理不了 {T} 个字"
|
| 54 |
+
x = self.tok(idx) + self.pos(np.arange(T))
|
| 55 |
+
else:
|
| 56 |
+
x = self.tok(idx) + Tensor(sinusoidal_encoding(T, self.tok.weight.shape[1]))
|
| 57 |
+
for b in self.blocks:
|
| 58 |
+
x = b(x)
|
| 59 |
+
return self.head(self.ln_f(x))
|
| 60 |
+
|
| 61 |
+
def state_floats(self, T):
|
| 62 |
+
"""推理到第 T 个字时 KV 缓存的大小(浮点数个数),随 T 线性增长。"""
|
| 63 |
+
return sum(2 * T * b.ln1.weight.shape[0] for b in self.blocks)
|
models/zeropoint.py
ADDED
|
@@ -0,0 +1,217 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ZeroPoint:分层记忆槽网络。
|
| 2 |
+
|
| 3 |
+
设计文档见 docs/ZeroPoint.md。核心是 MemoryMixer(代替 Transformer 里的自注意力):
|
| 4 |
+
|
| 5 |
+
每层有 H 个头,每个头有 M 个记忆槽 S[m](D 维向量),每个槽有自己的保留率 λ[m]。
|
| 6 |
+
读到第 t 个字时:
|
| 7 |
+
w_t = softmax(x_t Ww) * sigmoid(x_t Wg) 写地址(槽之间竞争)× 写强度
|
| 8 |
+
S[m] ← λ[m] · S[m] + w_t[m] · v_t 所有槽先按自己的速度遗忘,再写入
|
| 9 |
+
a_t = softmax(x_t Wr) 读地址
|
| 10 |
+
r_t = Σ_m a_t[m] · S[m] 从槽里读出
|
| 11 |
+
|
| 12 |
+
两种等价算法:
|
| 13 |
+
forward() 分块并行,训练用。序列切成每块 CHUNK 个字:
|
| 14 |
+
块内 r = (A ⊙ 因果掩码) @ v,A[i,j] = Σ_m a_i[m] w_j[m] λ[m]^(i-j)
|
| 15 |
+
利用 λ^(i-j) = λ^i · λ^(-j),一次矩阵乘
|
| 16 |
+
跨块 r += (a ⊙ λ^(i+1)) @ S_in S_in = 进入这一块时的槽状态
|
| 17 |
+
块尾 S_out = λ^c ⊙ S_in + (w ⊙ λ^(c-1-j))ᵀ @ v
|
| 18 |
+
λ^(-j) 最大 2^(CHUNK-1),不会溢出,所以序列多长都行
|
| 19 |
+
step() 递推形式,推理用:每个字 O(H·M·D),状态大小固定,没有 KV 缓存
|
| 20 |
+
|
| 21 |
+
槽的初始状态 S0 是可学习参数(不是全 0),这样序列开头读出来的向量不会接近 0、被 LayerNorm 放大噪声。
|
| 22 |
+
|
| 23 |
+
ema=True(指数滑动平均写入):S[m] ← λ[m]·S[m] + (1-λ[m])·w_t[m]·v_t
|
| 24 |
+
槽的值永远落在"写进去过的值"的范围里,不会随序列变长而越积越大。
|
| 25 |
+
系统直觉模型要连续跑几天,用这个;实现上只是把写地址 w 乘上 (1-λ),并行和递推两种形式都不用改。
|
| 26 |
+
"""
|
| 27 |
+
import numpy as np
|
| 28 |
+
|
| 29 |
+
from zerocg import Tensor, concat
|
| 30 |
+
from zerocg import functional as F
|
| 31 |
+
from zerocg import nn
|
| 32 |
+
|
| 33 |
+
from .transformer import sinusoidal_encoding
|
| 34 |
+
|
| 35 |
+
LAMBDA_MIN = 0.5 # λ = 0.5 + 0.5·sigmoid(θ) ∈ (0.5, 1)
|
| 36 |
+
CHUNK = 32 # 分块并行的块长:块内 λ^(-j) 最大 2^31,float32 绰绰有余
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def _layernorm_np(x, ln):
|
| 40 |
+
mu = x.mean(-1, keepdims=True)
|
| 41 |
+
var = ((x - mu) ** 2).mean(-1, keepdims=True)
|
| 42 |
+
return (x - mu) / np.sqrt(var + ln.eps) * ln.weight.data + ln.bias.data
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def _softmax_np(x):
|
| 46 |
+
e = np.exp(x - x.max(-1, keepdims=True))
|
| 47 |
+
return e / e.sum(-1, keepdims=True)
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def _gelu_np(a):
|
| 51 |
+
return 0.5 * a * (1 + np.tanh(np.sqrt(2 / np.pi) * (a + 0.044715 * a ** 3)))
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
class MemoryMixer(nn.Module):
|
| 55 |
+
"""参数里的开关只给消融实验用(experiments/ablation.py),默认值就是 ZeroPoint 本身:
|
| 56 |
+
gate 写强度门
|
| 57 |
+
write "softmax":写入时槽之间竞争;"sigmoid":每个槽独立决定写不写
|
| 58 |
+
init "hier":半衰期按存储层级铺开;"uniform":所有槽一样
|
| 59 |
+
learn_lambda 保留率 λ 是否可学
|
| 60 |
+
ema 写入按 (1-λ) 缩放(指数滑动平均),长时间运行不漂移
|
| 61 |
+
"""
|
| 62 |
+
|
| 63 |
+
def __init__(self, n_embd, n_head, n_slot, max_half_life=64.0,
|
| 64 |
+
gate=True, write="softmax", init="hier", learn_lambda=True, ema=False):
|
| 65 |
+
assert n_embd % n_head == 0 and write in ("softmax", "sigmoid") and init in ("hier", "uniform")
|
| 66 |
+
self.H, self.M, self.D = n_head, n_slot, n_embd // n_head
|
| 67 |
+
self.write_mode = write
|
| 68 |
+
self.ema = ema
|
| 69 |
+
self.read = nn.Linear(n_embd, n_head * n_slot)
|
| 70 |
+
self.write = nn.Linear(n_embd, n_head * n_slot)
|
| 71 |
+
self.gate = nn.Linear(n_embd, n_head) if gate else None
|
| 72 |
+
self.value = nn.Linear(n_embd, n_embd)
|
| 73 |
+
self.out = nn.Linear(n_embd, n_embd)
|
| 74 |
+
self.norm = nn.LayerNorm(self.D) # 每个头读出来的向量各自归一化(不同槽的量级差别很大)
|
| 75 |
+
# 初始化成"存储层级":槽 0 的半衰期约 1.5 个字(像寄存器),最后一个约 max_half_life 个字(像内存)
|
| 76 |
+
if init == "hier":
|
| 77 |
+
half_life = np.geomspace(1.5, max_half_life, n_slot)
|
| 78 |
+
else: # 对照:所有槽都取层级的几何中点
|
| 79 |
+
half_life = np.full(n_slot, np.sqrt(1.5 * max_half_life))
|
| 80 |
+
lam = 0.5 ** (1 / half_life)
|
| 81 |
+
theta = np.tile(np.log((2 * lam - 1) / (2 - 2 * lam)), (n_head, 1)) # sigmoid 的反函数,(H, M)
|
| 82 |
+
self.theta = nn.Parameter(theta) if learn_lambda else Tensor(theta)
|
| 83 |
+
self.s0 = nn.Parameter(np.random.randn(n_head, n_slot, self.D) * 0.1) # 可学习的初始状态
|
| 84 |
+
self.usage = None # 最近一次 forward 的 (读地址, 写地址) 平均值,各 (H, M),用来查槽有没有塌缩
|
| 85 |
+
|
| 86 |
+
def retention(self):
|
| 87 |
+
"""每个槽的保留率 λ,形状 (H, M)。"""
|
| 88 |
+
return F.sigmoid(self.theta) * (1 - LAMBDA_MIN) + LAMBDA_MIN
|
| 89 |
+
|
| 90 |
+
def forward(self, x):
|
| 91 |
+
B, T, C = x.shape
|
| 92 |
+
H, M, D = self.H, self.M, self.D
|
| 93 |
+
a = F.softmax(self.read(x).reshape(B, T, H, M), axis=-1).transpose(0, 2, 1, 3) # (B,H,T,M)
|
| 94 |
+
w = self.write(x).reshape(B, T, H, M)
|
| 95 |
+
w = F.softmax(w, axis=-1) if self.write_mode == "softmax" else F.sigmoid(w)
|
| 96 |
+
if self.gate is not None:
|
| 97 |
+
w = w * F.sigmoid(self.gate(x)).reshape(B, T, H, 1)
|
| 98 |
+
w = w.transpose(0, 2, 1, 3) # (B,H,T,M)
|
| 99 |
+
v = self.value(x).reshape(B, T, H, D).transpose(0, 2, 1, 3) # (B,H,T,D)
|
| 100 |
+
self.usage = (a.data.mean(axis=(0, 2)), w.data.mean(axis=(0, 2)))
|
| 101 |
+
|
| 102 |
+
lam = self.retention() # (H, M)
|
| 103 |
+
if self.ema:
|
| 104 |
+
w = w * (1 - lam).reshape(1, H, 1, M)
|
| 105 |
+
log_lam = lam.log().reshape(H, 1, M)
|
| 106 |
+
S = self.s0.reshape(1, H, M, D) # 进入第一块时的状态(对 batch 广播)
|
| 107 |
+
outs = []
|
| 108 |
+
for s in range(0, T, CHUNK):
|
| 109 |
+
e = min(s + CHUNK, T)
|
| 110 |
+
c = e - s
|
| 111 |
+
ac, wc, vc = a[:, :, s:e], w[:, :, s:e], v[:, :, s:e]
|
| 112 |
+
i = Tensor(np.arange(c, dtype=np.float64).reshape(1, c, 1))
|
| 113 |
+
decay = (log_lam * i).exp() # λ^i, (H,c,M)
|
| 114 |
+
A = (ac * decay) @ (wc * (log_lam * -i).exp()).transpose(0, 1, 3, 2) # (B,H,c,c)
|
| 115 |
+
A = A.masked_fill(np.triu(np.ones((c, c), dtype=bool), 1), 0.0)
|
| 116 |
+
outs.append(A @ vc + (ac * decay * lam.reshape(H, 1, M)) @ S) # 块内 + 跨块
|
| 117 |
+
if e < T:
|
| 118 |
+
tail = (log_lam * Tensor(np.arange(c - 1, -1, -1, dtype=np.float64).reshape(1, c, 1))).exp()
|
| 119 |
+
S = (lam ** c).reshape(1, H, M, 1) * S + (wc * tail).transpose(0, 1, 3, 2) @ vc
|
| 120 |
+
r = outs[0] if len(outs) == 1 else concat(outs, axis=2)
|
| 121 |
+
return self.out(self.norm(r).transpose(0, 2, 1, 3).reshape(B, T, C))
|
| 122 |
+
|
| 123 |
+
# ---------- 递推形式(numpy,无梯度)----------
|
| 124 |
+
def init_state(self, B):
|
| 125 |
+
return np.broadcast_to(self.s0.data, (B, self.H, self.M, self.D)).copy()
|
| 126 |
+
|
| 127 |
+
def step(self, x, S):
|
| 128 |
+
"""x: (B, C) 一个字的输入;S: (B, H, M, D) 记忆槽。返回 (输出, 新状态)。"""
|
| 129 |
+
B = x.shape[0]
|
| 130 |
+
H, M, D = self.H, self.M, self.D
|
| 131 |
+
lin = lambda layer, z: z @ layer.weight.data + layer.bias.data # noqa: E731
|
| 132 |
+
a = _softmax_np(lin(self.read, x).reshape(B, H, M))
|
| 133 |
+
w = lin(self.write, x).reshape(B, H, M)
|
| 134 |
+
w = _softmax_np(w) if self.write_mode == "softmax" else F.sigmoid_np(w)
|
| 135 |
+
if self.gate is not None:
|
| 136 |
+
w = w * F.sigmoid_np(lin(self.gate, x))[:, :, None]
|
| 137 |
+
v = lin(self.value, x).reshape(B, H, D)
|
| 138 |
+
lam = self.retention().data
|
| 139 |
+
if self.ema:
|
| 140 |
+
w = w * (1 - lam)[None]
|
| 141 |
+
self.last_addr = (a, w) # 诊断用:这一步的读地址和写地址,(B, H, M)
|
| 142 |
+
S = lam[None, :, :, None] * S + w[..., None] * v[:, :, None, :]
|
| 143 |
+
r = np.einsum("bhm,bhmd->bhd", a, S)
|
| 144 |
+
r = _layernorm_np(r, self.norm).reshape(B, H * D)
|
| 145 |
+
return lin(self.out, r), S
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
class Block(nn.Module):
|
| 149 |
+
def __init__(self, n_embd, n_head, n_slot, dropout=0.0, mlp_hidden=None, **mixer_opts):
|
| 150 |
+
hidden = mlp_hidden or 4 * n_embd
|
| 151 |
+
self.ln1 = nn.LayerNorm(n_embd)
|
| 152 |
+
self.mix = MemoryMixer(n_embd, n_head, n_slot, **mixer_opts)
|
| 153 |
+
self.ln2 = nn.LayerNorm(n_embd)
|
| 154 |
+
self.fc1 = nn.Linear(n_embd, hidden)
|
| 155 |
+
self.fc2 = nn.Linear(hidden, n_embd)
|
| 156 |
+
self.drop = nn.Dropout(dropout)
|
| 157 |
+
|
| 158 |
+
def forward(self, x):
|
| 159 |
+
x = x + self.drop(self.mix(self.ln1(x)))
|
| 160 |
+
return x + self.drop(self.fc2(F.gelu(self.fc1(self.ln2(x)))))
|
| 161 |
+
|
| 162 |
+
def step(self, x, S):
|
| 163 |
+
h, S = self.mix.step(_layernorm_np(x, self.ln1), S)
|
| 164 |
+
x = x + h
|
| 165 |
+
h = _gelu_np(_layernorm_np(x, self.ln2) @ self.fc1.weight.data + self.fc1.bias.data)
|
| 166 |
+
return x + h @ self.fc2.weight.data + self.fc2.bias.data, S
|
| 167 |
+
|
| 168 |
+
|
| 169 |
+
class ZeroX(nn.Module):
|
| 170 |
+
"""字符级语言模型。默认没有位置编码:位置信息来自各个槽不同的遗忘速度。
|
| 171 |
+
pos="sin" 只给消融实验用:额外加上论文的正弦位置编码,看看还有没有用。
|
| 172 |
+
mlp_hidden 默认 4 × n_embd;消融实验用它把各变体的参数量对齐。"""
|
| 173 |
+
|
| 174 |
+
def __init__(self, vocab_size, n_embd=64, n_head=4, n_slot=16, n_layer=2, dropout=0.0,
|
| 175 |
+
pos="none", mlp_hidden=None, **mixer_opts):
|
| 176 |
+
assert pos in ("none", "sin")
|
| 177 |
+
self.pos_kind = pos
|
| 178 |
+
self.tok = nn.Embedding(vocab_size, n_embd)
|
| 179 |
+
self.blocks = [Block(n_embd, n_head, n_slot, dropout, mlp_hidden, **mixer_opts) for _ in range(n_layer)]
|
| 180 |
+
self.ln_f = nn.LayerNorm(n_embd)
|
| 181 |
+
self.head = nn.Linear(n_embd, vocab_size)
|
| 182 |
+
|
| 183 |
+
def forward(self, idx):
|
| 184 |
+
x = self.tok(idx)
|
| 185 |
+
if self.pos_kind == "sin":
|
| 186 |
+
x = x + Tensor(sinusoidal_encoding(idx.shape[1], x.shape[-1]))
|
| 187 |
+
for b in self.blocks:
|
| 188 |
+
x = b(x)
|
| 189 |
+
return self.head(self.ln_f(x))
|
| 190 |
+
|
| 191 |
+
def init_state(self, B):
|
| 192 |
+
return [b.mix.init_state(B) for b in self.blocks]
|
| 193 |
+
|
| 194 |
+
def step(self, idx_t, states, t=0):
|
| 195 |
+
"""idx_t: (B,) 第 t 个字。返回 (logits (B, V), 新状态)。"""
|
| 196 |
+
x = self.tok.weight.data[idx_t]
|
| 197 |
+
if self.pos_kind == "sin":
|
| 198 |
+
x = x + sinusoidal_encoding(t + 1, x.shape[-1])[t]
|
| 199 |
+
new = []
|
| 200 |
+
for b, S in zip(self.blocks, states):
|
| 201 |
+
x, S = b.step(x, S)
|
| 202 |
+
new.append(S)
|
| 203 |
+
return _layernorm_np(x, self.ln_f) @ self.head.weight.data + self.head.bias.data, new
|
| 204 |
+
|
| 205 |
+
def forward_recurrent(self, idx):
|
| 206 |
+
"""逐字递推跑完整个序列,返回 (B, T, V) 的 logits。"""
|
| 207 |
+
idx = np.asarray(idx)
|
| 208 |
+
states = self.init_state(idx.shape[0])
|
| 209 |
+
out = []
|
| 210 |
+
for t in range(idx.shape[1]):
|
| 211 |
+
logits, states = self.step(idx[:, t], states, t)
|
| 212 |
+
out.append(logits)
|
| 213 |
+
return np.stack(out, axis=1)
|
| 214 |
+
|
| 215 |
+
def state_floats(self):
|
| 216 |
+
"""推理时需要保存的状态大小(浮点数个数),和序列长度无关。"""
|
| 217 |
+
return sum(b.mix.H * b.mix.M * b.mix.D for b in self.blocks)
|
requirements.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
numpy
|
scripts/upload.md
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Publishing ZeroPoint to the Hugging Face Hub
|
| 2 |
+
|
| 3 |
+
This note records the exact steps used to publish this release tree as a **public** model
|
| 4 |
+
repository. It is for the maintainer performing the upload; users of the release do not
|
| 5 |
+
need it. Never commit a real access token, and never print one in a terminal that records
|
| 6 |
+
its output.
|
| 7 |
+
|
| 8 |
+
The commands below were checked against `huggingface_hub` 1.27.0 (`hf` CLI).
|
| 9 |
+
|
| 10 |
+
## 0. Environment: the local proxy is mandatory here
|
| 11 |
+
|
| 12 |
+
In this environment `huggingface.co` is only reachable through the local HTTP proxy at
|
| 13 |
+
`http://127.0.0.1:7897`. Set both proxy variables in every shell that talks to the Hub,
|
| 14 |
+
and pin the official endpoint:
|
| 15 |
+
|
| 16 |
+
```bash
|
| 17 |
+
export HTTPS_PROXY=http://127.0.0.1:7897
|
| 18 |
+
export HTTP_PROXY=http://127.0.0.1:7897
|
| 19 |
+
export HF_ENDPOINT=https://huggingface.co
|
| 20 |
+
```
|
| 21 |
+
|
| 22 |
+
If the proxy is not running on port 7897, the commands below hang and then fail with a
|
| 23 |
+
connection or SSL error — start the proxy first. Do not leave a mirror or an offline
|
| 24 |
+
`HF_ENDPOINT` in the environment.
|
| 25 |
+
|
| 26 |
+
## 1. Authenticate (token required)
|
| 27 |
+
|
| 28 |
+
Create a **write** token at <https://huggingface.co/settings/tokens>. Then either log in
|
| 29 |
+
interactively (the token is not echoed):
|
| 30 |
+
|
| 31 |
+
```bash
|
| 32 |
+
hf auth login
|
| 33 |
+
```
|
| 34 |
+
|
| 35 |
+
or pass the token from an environment variable, without ever writing it to a file:
|
| 36 |
+
|
| 37 |
+
```bash
|
| 38 |
+
read -rs HF_TOKEN && export HF_TOKEN # paste at the prompt; input is hidden
|
| 39 |
+
hf auth login --token "$HF_TOKEN"
|
| 40 |
+
```
|
| 41 |
+
|
| 42 |
+
`hf auth login --token "$HF_TOKEN"` stores the token in the local Hugging Face config
|
| 43 |
+
(`~/.cache/huggingface/`), so the later commands pick it up automatically. Verify:
|
| 44 |
+
|
| 45 |
+
```bash
|
| 46 |
+
hf auth whoami
|
| 47 |
+
```
|
| 48 |
+
|
| 49 |
+
## 2. Create the public model repository
|
| 50 |
+
|
| 51 |
+
```bash
|
| 52 |
+
hf repos create <your-username>/ZeroPoint --type model --public
|
| 53 |
+
```
|
| 54 |
+
|
| 55 |
+
Replace `<your-username>` with the account (or organization) that owns the release. If the
|
| 56 |
+
repository already exists, skip this step and make sure its visibility is public in the
|
| 57 |
+
repository settings.
|
| 58 |
+
|
| 59 |
+
## 3. Upload the release tree
|
| 60 |
+
|
| 61 |
+
Upload the **contents** of `repo/` so that `README.md`, `LICENSE`, `docs/`, `models/`,
|
| 62 |
+
`experiments/`, `tests/`, etc. land at the repository root:
|
| 63 |
+
|
| 64 |
+
```bash
|
| 65 |
+
cd "<release-root>" # the directory containing README.md, LICENSE, models/, docs/ (i.e. the parent of this scripts/ folder)
|
| 66 |
+
|
| 67 |
+
hf upload <your-username>/ZeroPoint . --type model \
|
| 68 |
+
--commit-message "ZeroPoint v0 research prototype: architecture, ZeroCG framework, comparison experiments"
|
| 69 |
+
```
|
| 70 |
+
|
| 71 |
+
This is a single-commit upload of the whole tree. It includes
|
| 72 |
+
`tests/golden/zeropoint_ref.npz` (the fixed reference data required by the self-checks)
|
| 73 |
+
and must not include model weights, training runs, corpora, or tokens — see
|
| 74 |
+
`.gitignore`; the release tree itself contains none of those.
|
| 75 |
+
|
| 76 |
+
If the upload is interrupted, re-run the same command: files already uploaded are skipped
|
| 77 |
+
and the commit is retried.
|
| 78 |
+
|
| 79 |
+
## 4. Verify the published repository
|
| 80 |
+
|
| 81 |
+
```bash
|
| 82 |
+
hf repos ls
|
| 83 |
+
hf auth whoami
|
| 84 |
+
|
| 85 |
+
# Confirm the public API view, through the proxy:
|
| 86 |
+
curl -sS --proxy http://127.0.0.1:7897 \
|
| 87 |
+
"https://huggingface.co/api/models/<your-username>/ZeroPoint" | head -c 2000
|
| 88 |
+
```
|
| 89 |
+
|
| 90 |
+
Then open <https://huggingface.co/<your-username>/ZeroPoint> in a browser and confirm
|
| 91 |
+
that `README.md`, `LICENSE`, `docs/ZeroPoint.md`, `docs/ZeroPoint_EN.md`, and the
|
| 92 |
+
`tests/golden/` file are all present.
|
tests/golden/zeropoint_ref.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5a2b72d72aeac2cd51bc8987ba5e3abbc29ebf623e2f93f40ac5995d77336ca0
|
| 3 |
+
size 85168
|
tests/make_golden.py
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""生成 ZeroPoint 的参考数据 tests/golden/zeropoint_ref.npz:固定的参数 + 输入 + 标准输出。
|
| 2 |
+
|
| 3 |
+
所有实现(zerocg 并行形式、numpy 递推形式、以后内核里的 C 版本)都拿它对答案,
|
| 4 |
+
谁和它对不上谁就有 bug。**只有在有意修改架构时才重新生成**,并在提交说明里写清楚为什么。
|
| 5 |
+
|
| 6 |
+
python3 tests/make_golden.py
|
| 7 |
+
"""
|
| 8 |
+
import os
|
| 9 |
+
import sys
|
| 10 |
+
|
| 11 |
+
import numpy as np
|
| 12 |
+
|
| 13 |
+
HERE = os.path.dirname(os.path.abspath(__file__))
|
| 14 |
+
sys.path.insert(0, os.path.join(HERE, ".."))
|
| 15 |
+
import zerocg as zc # noqa: E402
|
| 16 |
+
from models import ZeroX # noqa: E402
|
| 17 |
+
|
| 18 |
+
CONFIG = dict(vocab_size=13, n_embd=16, n_head=2, n_slot=4, n_layer=2)
|
| 19 |
+
|
| 20 |
+
if __name__ == "__main__":
|
| 21 |
+
zc.set_default_dtype(np.float64)
|
| 22 |
+
np.random.seed(2026)
|
| 23 |
+
model = ZeroX(**CONFIG).eval()
|
| 24 |
+
# 参数不用初始值:随机扰动一下,避免初始化时的对称性掩盖 bug
|
| 25 |
+
for p in model.parameters():
|
| 26 |
+
p.data = p.data + np.random.randn(*p.shape) * 0.3
|
| 27 |
+
idx = np.random.randint(0, CONFIG["vocab_size"], size=(2, 70)) # 跨 3 个块
|
| 28 |
+
logits = model(idx).data
|
| 29 |
+
out = os.path.join(HERE, "golden", "zeropoint_ref.npz")
|
| 30 |
+
np.savez(out, idx=idx, logits=logits, **{"param:" + k: v for k, v in model.state_dict().items()},
|
| 31 |
+
**{"config:" + k: v for k, v in CONFIG.items()})
|
| 32 |
+
print(f"写入 {out}:输入 {idx.shape},输出 {logits.shape},{len(model.state_dict())} 个参数张量")
|
tests/test_zerocg.py
ADDED
|
@@ -0,0 +1,250 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""框架自检:用数值差分验证每个算子 / 层的反向传播,并跑两个小模型确认能学会。
|
| 2 |
+
运行: python3 tests/test_zerocg.py
|
| 3 |
+
"""
|
| 4 |
+
import os
|
| 5 |
+
import sys
|
| 6 |
+
|
| 7 |
+
import numpy as np
|
| 8 |
+
|
| 9 |
+
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
|
| 10 |
+
import zerocg as zc # noqa: E402
|
| 11 |
+
from zerocg import functional as F # noqa: E402
|
| 12 |
+
from zerocg import nn # noqa: E402
|
| 13 |
+
|
| 14 |
+
zc.set_default_dtype(np.float64) # 数值梯度检查需要 float64 的精度
|
| 15 |
+
rng = np.random.default_rng(0)
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def grad_check(fn, *shapes, tol=1e-5, positive=False):
|
| 19 |
+
"""fn(*tensors) -> Tensor。比较自动求导和中心差分的梯度(float64)。"""
|
| 20 |
+
arrays = [rng.standard_normal(s) for s in shapes]
|
| 21 |
+
if positive:
|
| 22 |
+
arrays = [np.abs(a) + 0.5 for a in arrays]
|
| 23 |
+
ts = [zc.Tensor(a.copy(), requires_grad=True) for a in arrays]
|
| 24 |
+
w = np.asarray(rng.standard_normal(fn(*ts).shape)) # 随机投影成标量,检查完整雅可比
|
| 25 |
+
(fn(*ts) * zc.Tensor(w)).sum().backward()
|
| 26 |
+
eps = 1e-6
|
| 27 |
+
for t, a in zip(ts, arrays):
|
| 28 |
+
num = np.zeros_like(a)
|
| 29 |
+
for i in np.ndindex(a.shape):
|
| 30 |
+
for sign in (1, -1):
|
| 31 |
+
b = a.copy()
|
| 32 |
+
b[i] += sign * eps
|
| 33 |
+
inputs = [zc.Tensor(b) if x is t else zc.Tensor(x.data) for x in ts]
|
| 34 |
+
with zc.no_grad():
|
| 35 |
+
num[i] += sign * (fn(*inputs).data * w).sum()
|
| 36 |
+
num[i] /= 2 * eps
|
| 37 |
+
err = np.abs(num - t.grad).max() / max(1.0, np.abs(num).max())
|
| 38 |
+
assert err < tol, f"梯度误差 {err:.2e}"
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def module_grad_check(module, x):
|
| 42 |
+
for p in module.parameters():
|
| 43 |
+
p.data = p.data.astype(np.float64)
|
| 44 |
+
params = module.parameters()
|
| 45 |
+
# 直接对参数逐个做差分
|
| 46 |
+
xt = zc.Tensor(x)
|
| 47 |
+
w = np.asarray(rng.standard_normal(module(xt).shape))
|
| 48 |
+
for p in params:
|
| 49 |
+
p.grad = None
|
| 50 |
+
(module(xt) * zc.Tensor(w)).sum().backward()
|
| 51 |
+
eps = 1e-6
|
| 52 |
+
for p in params:
|
| 53 |
+
num = np.zeros_like(p.data)
|
| 54 |
+
for i in np.ndindex(p.shape):
|
| 55 |
+
old = p.data[i]
|
| 56 |
+
p.data[i] = old + eps
|
| 57 |
+
with zc.no_grad():
|
| 58 |
+
up = (module(xt).data * w).sum()
|
| 59 |
+
p.data[i] = old - eps
|
| 60 |
+
with zc.no_grad():
|
| 61 |
+
down = (module(xt).data * w).sum()
|
| 62 |
+
p.data[i] = old
|
| 63 |
+
num[i] = (up - down) / (2 * eps)
|
| 64 |
+
err = np.abs(num - p.grad).max() / max(1.0, np.abs(num).max())
|
| 65 |
+
assert err < 1e-5, f"参数梯度误差 {err:.2e}"
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# ---------------- 算子 ----------------
|
| 69 |
+
def test_ops():
|
| 70 |
+
grad_check(lambda a, b: a + b, (3, 4), (4,))
|
| 71 |
+
grad_check(lambda a, b: a - b, (3, 4), (3, 1))
|
| 72 |
+
grad_check(lambda a, b: a * b, (2, 3), (2, 3))
|
| 73 |
+
grad_check(lambda a, b: a / b, (2, 3), (3,), positive=True)
|
| 74 |
+
grad_check(lambda a: a ** 3, (4,))
|
| 75 |
+
grad_check(lambda a: a.sqrt(), (4,), positive=True)
|
| 76 |
+
grad_check(lambda a: -a + 2.0 - a * 3.0, (4,))
|
| 77 |
+
grad_check(lambda a: 1.0 / a, (4,), positive=True)
|
| 78 |
+
grad_check(lambda a: a.exp(), (3, 2))
|
| 79 |
+
grad_check(lambda a: a.log(), (3, 2), positive=True)
|
| 80 |
+
grad_check(lambda a: a.tanh(), (5,))
|
| 81 |
+
grad_check(lambda a: a.relu(), (5,))
|
| 82 |
+
grad_check(lambda a: a.masked_fill(np.eye(3, dtype=bool), -5.0), (3, 3))
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def test_reductions_and_shapes():
|
| 86 |
+
grad_check(lambda a: a.sum(), (3, 4))
|
| 87 |
+
grad_check(lambda a: a.sum(axis=1), (3, 4))
|
| 88 |
+
grad_check(lambda a: a.mean(axis=0, keepdims=True), (3, 4))
|
| 89 |
+
grad_check(lambda a: a.max(axis=-1), (3, 4))
|
| 90 |
+
grad_check(lambda a: a.reshape(4, 3), (3, 4))
|
| 91 |
+
grad_check(lambda a: a.transpose(2, 0, 1), (2, 3, 4))
|
| 92 |
+
grad_check(lambda a: a[1:, ::2], (3, 4))
|
| 93 |
+
grad_check(lambda a: a[np.array([0, 2, 0])], (3, 4)) # 重复索引
|
| 94 |
+
grad_check(lambda a, b: zc.concat([a, b, a], axis=1), (2, 3), (2, 1))
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def test_matmul():
|
| 98 |
+
grad_check(lambda a, b: a @ b, (3, 4), (4, 5))
|
| 99 |
+
grad_check(lambda a, b: a @ b, (2, 3, 4), (4, 5)) # 批量 @ 共享权重
|
| 100 |
+
grad_check(lambda a, b: a @ b, (2, 2, 3, 4), (2, 2, 4, 3))
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def test_functional():
|
| 104 |
+
grad_check(lambda a: F.softmax(a, axis=-1), (3, 5))
|
| 105 |
+
grad_check(lambda a: F.log_softmax(a, axis=-1), (3, 5))
|
| 106 |
+
grad_check(lambda a: F.gelu(a), (6,))
|
| 107 |
+
grad_check(lambda a: F.sigmoid(a), (6,))
|
| 108 |
+
grad_check(lambda a: F.cross_entropy(a, np.array([1, 0, 4])), (3, 5))
|
| 109 |
+
# 和直接的数学定义对比数值
|
| 110 |
+
x = rng.standard_normal((4, 6))
|
| 111 |
+
t = np.array([0, 5, 2, 2])
|
| 112 |
+
ref = -np.log(np.exp(x) / np.exp(x).sum(1, keepdims=True))[np.arange(4), t].mean()
|
| 113 |
+
assert abs(F.cross_entropy(zc.Tensor(x), t).item() - ref) < 1e-10
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
def test_layers():
|
| 117 |
+
module_grad_check(nn.Linear(4, 3), rng.standard_normal((2, 4)))
|
| 118 |
+
module_grad_check(nn.LayerNorm(5), rng.standard_normal((3, 5)))
|
| 119 |
+
module_grad_check(nn.Sequential(nn.Linear(3, 4), nn.GELU(), nn.Linear(4, 2)), rng.standard_normal((2, 3)))
|
| 120 |
+
module_grad_check(nn.CausalSelfAttention(8, 2), rng.standard_normal((2, 4, 8)))
|
| 121 |
+
|
| 122 |
+
emb = nn.Embedding(5, 3)
|
| 123 |
+
out = emb(np.array([[1, 1, 4]]))
|
| 124 |
+
out.sum().backward()
|
| 125 |
+
assert emb.weight.grad[1].tolist() == [2, 2, 2] and emb.weight.grad[0].tolist() == [0, 0, 0]
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
def test_attention_is_causal():
|
| 129 |
+
attn = nn.CausalSelfAttention(8, 2).eval()
|
| 130 |
+
x = rng.standard_normal((1, 5, 8)).astype(np.float32)
|
| 131 |
+
y1 = attn(zc.Tensor(x)).data
|
| 132 |
+
x[0, 4] += 10 # 改最后一个位置,前面的输出不应变化
|
| 133 |
+
y2 = attn(zc.Tensor(x)).data
|
| 134 |
+
assert np.allclose(y1[0, :4], y2[0, :4]) and not np.allclose(y1[0, 4], y2[0, 4])
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def test_count_flops():
|
| 138 |
+
with zc.count_flops() as c:
|
| 139 |
+
zc.randn(2, 3, 4) @ zc.randn(4, 5)
|
| 140 |
+
assert c.flops == 2 * (2 * 3 * 5) * 4
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
def test_default_dtype():
|
| 144 |
+
zc.set_default_dtype(np.float32)
|
| 145 |
+
try:
|
| 146 |
+
x = zc.Tensor(np.ones(3)) # float64 数组进来也变成 float32
|
| 147 |
+
assert x.dtype == np.float32 and (x * 0.5).dtype == np.float32
|
| 148 |
+
assert nn.Linear(2, 2).weight.dtype == np.float32
|
| 149 |
+
assert zc.Tensor(np.arange(3)).dtype.kind == "i" # 整数下标不动
|
| 150 |
+
assert np.isfinite(F.sigmoid(zc.Tensor([-1000.0, 0.0, 1000.0])).data).all()
|
| 151 |
+
finally:
|
| 152 |
+
zc.set_default_dtype(np.float64)
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
def test_state_dict_roundtrip(tmp="/tmp"):
|
| 156 |
+
path = os.path.join(os.environ.get("TMPDIR", tmp), "zerocg_test.npz")
|
| 157 |
+
a = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4))
|
| 158 |
+
a.save(path)
|
| 159 |
+
b = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4)).load(path)
|
| 160 |
+
x = zc.randn(2, 3)
|
| 161 |
+
assert np.allclose(a(x).data, b(x).data)
|
| 162 |
+
assert [n for n, _ in a.named_parameters()] == ["layers.0.weight", "layers.0.bias", "layers.1.weight", "layers.1.bias"]
|
| 163 |
+
|
| 164 |
+
|
| 165 |
+
def test_optimizer_resume():
|
| 166 |
+
"""训练 10 步 == 训练 5 步、存档、新建模型和优化器、读档、再训练 5 步。"""
|
| 167 |
+
X, y = rng.standard_normal((8, 3)), np.array([0, 1, 1, 0, 1, 0, 0, 1])
|
| 168 |
+
|
| 169 |
+
def make():
|
| 170 |
+
np.random.seed(7)
|
| 171 |
+
m = nn.Sequential(nn.Linear(3, 4), nn.Tanh(), nn.Linear(4, 2))
|
| 172 |
+
return m, zc.optim.AdamW(m.parameters(), lr=0.05, weight_decay=0.01)
|
| 173 |
+
|
| 174 |
+
def train(m, opt, n):
|
| 175 |
+
for _ in range(n):
|
| 176 |
+
loss = F.cross_entropy(m(zc.Tensor(X)), y)
|
| 177 |
+
opt.zero_grad()
|
| 178 |
+
loss.backward()
|
| 179 |
+
opt.step()
|
| 180 |
+
|
| 181 |
+
a, opt_a = make()
|
| 182 |
+
train(a, opt_a, 10)
|
| 183 |
+
b, opt_b = make()
|
| 184 |
+
train(b, opt_b, 5)
|
| 185 |
+
model_state, opt_state = b.state_dict(), opt_b.state_dict()
|
| 186 |
+
c, opt_c = make()
|
| 187 |
+
c.load_state_dict(model_state)
|
| 188 |
+
opt_c.load_state_dict(opt_state)
|
| 189 |
+
train(c, opt_c, 5)
|
| 190 |
+
assert all(np.array_equal(p, q) for p, q in zip(a.state_dict().values(), c.state_dict().values()))
|
| 191 |
+
|
| 192 |
+
|
| 193 |
+
# ---------------- 端到端学习 ----------------
|
| 194 |
+
def test_learns_xor():
|
| 195 |
+
np.random.seed(0)
|
| 196 |
+
X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32)
|
| 197 |
+
y = np.array([0, 1, 1, 0])
|
| 198 |
+
model = nn.Sequential(nn.Linear(2, 16), nn.Tanh(), nn.Linear(16, 2))
|
| 199 |
+
opt = zc.optim.AdamW(model.parameters(), lr=0.05)
|
| 200 |
+
for _ in range(300):
|
| 201 |
+
loss = F.cross_entropy(model(zc.Tensor(X)), y)
|
| 202 |
+
opt.zero_grad()
|
| 203 |
+
loss.backward()
|
| 204 |
+
opt.step()
|
| 205 |
+
assert (model(zc.Tensor(X)).data.argmax(1) == y).all(), f"XOR 没学会, loss={loss.item()}"
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
def test_attention_memorizes_sequence():
|
| 209 |
+
"""一层注意力 + 嵌入,记住一个短序列的下一个字符。"""
|
| 210 |
+
np.random.seed(0)
|
| 211 |
+
seq = np.array([0, 1, 2, 3, 2, 1, 0, 3])
|
| 212 |
+
x, y = seq[None, :-1], seq[1:]
|
| 213 |
+
|
| 214 |
+
class Tiny(nn.Module):
|
| 215 |
+
def __init__(self):
|
| 216 |
+
self.tok = nn.Embedding(4, 16)
|
| 217 |
+
self.pos = nn.Embedding(8, 16)
|
| 218 |
+
self.attn = nn.CausalSelfAttention(16, 2)
|
| 219 |
+
self.ln = nn.LayerNorm(16)
|
| 220 |
+
self.head = nn.Linear(16, 4)
|
| 221 |
+
|
| 222 |
+
def forward(self, idx):
|
| 223 |
+
T = idx.shape[1]
|
| 224 |
+
h = self.tok(idx) + self.pos(np.arange(T))
|
| 225 |
+
h = h + self.attn(self.ln(h))
|
| 226 |
+
return self.head(h)
|
| 227 |
+
|
| 228 |
+
model = Tiny()
|
| 229 |
+
opt = zc.optim.AdamW(model.parameters(), lr=0.02)
|
| 230 |
+
for _ in range(200):
|
| 231 |
+
loss = F.cross_entropy(model(x).reshape(-1, 4), y)
|
| 232 |
+
opt.zero_grad()
|
| 233 |
+
loss.backward()
|
| 234 |
+
zc.optim.clip_grad_norm(model.parameters(), 1.0)
|
| 235 |
+
opt.step()
|
| 236 |
+
assert loss.item() < 0.05, f"loss={loss.item()}"
|
| 237 |
+
|
| 238 |
+
|
| 239 |
+
if __name__ == "__main__":
|
| 240 |
+
tests = [(k, v) for k, v in list(globals().items()) if k.startswith("test_")]
|
| 241 |
+
failed = 0
|
| 242 |
+
for name, fn in tests:
|
| 243 |
+
try:
|
| 244 |
+
fn()
|
| 245 |
+
print(f" ok {name}")
|
| 246 |
+
except Exception as e: # noqa: BLE001
|
| 247 |
+
failed += 1
|
| 248 |
+
print(f" FAIL {name}: {type(e).__name__}: {e}")
|
| 249 |
+
print(f"\n{len(tests) - failed}/{len(tests)} 通过")
|
| 250 |
+
sys.exit(1 if failed else 0)
|
tests/test_zeropoint.py
ADDED
|
@@ -0,0 +1,153 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ZeroPoint 架构自检。运行: python3 tests/test_zeropoint.py"""
|
| 2 |
+
import itertools
|
| 3 |
+
import os
|
| 4 |
+
import sys
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
|
| 8 |
+
HERE = os.path.dirname(os.path.abspath(__file__))
|
| 9 |
+
sys.path.insert(0, os.path.join(HERE, ".."))
|
| 10 |
+
sys.path.insert(0, HERE)
|
| 11 |
+
import zerocg as zc # noqa: E402
|
| 12 |
+
from models import ZeroX # noqa: E402
|
| 13 |
+
from models.zeropoint import CHUNK, MemoryMixer # noqa: E402
|
| 14 |
+
from test_zerocg import module_grad_check # noqa: E402
|
| 15 |
+
|
| 16 |
+
zc.set_default_dtype(np.float64)
|
| 17 |
+
rng = np.random.default_rng(0)
|
| 18 |
+
|
| 19 |
+
# 所有开关的组合(消融实验会用到的每一种都要验证两种算法一致)
|
| 20 |
+
OPTIONS = [dict(zip(("gate", "write", "init", "learn_lambda", "pos", "ema"), combo))
|
| 21 |
+
for combo in itertools.product((True, False), ("softmax", "sigmoid"), ("hier", "uniform"),
|
| 22 |
+
(True, False), ("none", "sin"), (False, True))]
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def test_parallel_equals_recurrent():
|
| 26 |
+
"""训练用的分块并行形式和推理用的递推形式,在所有开关组合下都必须一致。
|
| 27 |
+
长度 300:跨好几个块,也远超旧版并行形式 120 字的上限。"""
|
| 28 |
+
idx = rng.integers(0, 11, size=(2, 300))
|
| 29 |
+
worst = 0.0
|
| 30 |
+
for opts in OPTIONS:
|
| 31 |
+
np.random.seed(1)
|
| 32 |
+
model = ZeroX(vocab_size=11, n_embd=16, n_head=2, n_slot=5, n_layer=2, **opts).eval()
|
| 33 |
+
worst = max(worst, np.abs(model(idx).data - model.forward_recurrent(idx)).max())
|
| 34 |
+
assert worst < 1e-8, worst
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def test_matches_golden_reference():
|
| 38 |
+
"""和 tests/golden/zeropoint_ref.npz 对答案:防止以后改代码时并行或递推形式悄悄变了。"""
|
| 39 |
+
ref = np.load(os.path.join(HERE, "golden", "zeropoint_ref.npz"))
|
| 40 |
+
cfg = {k[7:]: int(ref[k]) for k in ref.files if k.startswith("config:")}
|
| 41 |
+
model = ZeroX(**cfg).eval()
|
| 42 |
+
model.load_state_dict({k[6:]: ref[k] for k in ref.files if k.startswith("param:")})
|
| 43 |
+
assert np.abs(model(ref["idx"]).data - ref["logits"]).max() < 1e-10
|
| 44 |
+
assert np.abs(model.forward_recurrent(ref["idx"]) - ref["logits"]).max() < 1e-8
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def test_causal_across_chunks():
|
| 48 |
+
np.random.seed(2)
|
| 49 |
+
model = ZeroX(vocab_size=7, n_embd=16, n_head=2, n_slot=4).eval()
|
| 50 |
+
idx = rng.integers(0, 7, size=(1, 3 * CHUNK))
|
| 51 |
+
y1 = model(idx).data
|
| 52 |
+
p = CHUNK + 5 # 改第二块里的一个字
|
| 53 |
+
idx[0, p] = (idx[0, p] + 1) % 7
|
| 54 |
+
y2 = model(idx).data
|
| 55 |
+
assert np.allclose(y1[0, :p], y2[0, :p]) and not np.allclose(y1[0, p:], y2[0, p:])
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def test_ema_state_is_bounded():
|
| 59 |
+
"""ema=True:写入恒为常数 v 时,槽会收敛到 v 本身,不会随长度无限变大(ema=False 时会涨到约 v/(1-λ))。"""
|
| 60 |
+
np.random.seed(6)
|
| 61 |
+
for ema, bounded in ((True, True), (False, False)):
|
| 62 |
+
mixer = MemoryMixer(8, 1, 4, ema=ema)
|
| 63 |
+
S = mixer.init_state(1) * 0
|
| 64 |
+
x = np.ones((1, 8))
|
| 65 |
+
for _ in range(3000):
|
| 66 |
+
_, S = mixer.step(x, S)
|
| 67 |
+
v = np.abs(x @ mixer.value.weight.data + mixer.value.bias.data).max()
|
| 68 |
+
assert (np.abs(S).max() <= v * 1.0001) == bounded, (ema, np.abs(S).max(), v)
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def test_mixer_gradients():
|
| 72 |
+
np.random.seed(3)
|
| 73 |
+
module_grad_check(MemoryMixer(8, 2, 3), rng.standard_normal((2, 6, 8)))
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def test_mixer_gradients_across_chunks():
|
| 77 |
+
"""梯度要能穿过块与块之间传递的状态,也要传到可学习的初始状态 s0。"""
|
| 78 |
+
np.random.seed(3)
|
| 79 |
+
mixer = MemoryMixer(4, 1, 2)
|
| 80 |
+
module_grad_check(mixer, rng.standard_normal((1, CHUNK + 3, 4)))
|
| 81 |
+
assert np.abs(mixer.s0.grad).max() > 0
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def test_initial_state_is_learned():
|
| 85 |
+
model = ZeroX(vocab_size=5, n_embd=8, n_head=2, n_slot=3, n_layer=1)
|
| 86 |
+
assert np.allclose(model.init_state(3)[0][2], model.blocks[0].mix.s0.data)
|
| 87 |
+
assert "blocks.0.mix.s0" in dict(model.named_parameters())
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def test_retention_init_is_hierarchical():
|
| 91 |
+
lam = MemoryMixer(16, 2, 6, max_half_life=64).retention().data
|
| 92 |
+
half_life = np.log(0.5) / np.log(lam[0])
|
| 93 |
+
assert np.all(np.diff(lam[0]) > 0) and abs(half_life[0] - 1.5) < 1e-3 and abs(half_life[-1] - 64) < 1e-2
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
def test_usage_recorded():
|
| 97 |
+
model = ZeroX(vocab_size=5, n_embd=8, n_head=2, n_slot=3, n_layer=1)
|
| 98 |
+
model(rng.integers(0, 5, size=(2, 10)))
|
| 99 |
+
read, write = model.blocks[0].mix.usage
|
| 100 |
+
assert read.shape == write.shape == (2, 3) and np.allclose(read.sum(-1), 1)
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def test_state_is_constant_size():
|
| 104 |
+
model = ZeroX(vocab_size=5, n_embd=16, n_head=2, n_slot=4, n_layer=3)
|
| 105 |
+
states = model.init_state(1)
|
| 106 |
+
for t in range(300):
|
| 107 |
+
_, states = model.step(np.array([t % 5]), states, t)
|
| 108 |
+
assert sum(s.size for s in states) == model.state_floats() == 3 * 2 * 4 * 8
|
| 109 |
+
assert all(np.isfinite(s).all() for s in states)
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def test_float32_long_sequence_is_finite():
|
| 113 |
+
"""float32 下 1000 个字的分块并行不溢出,并且和递推形式一致。"""
|
| 114 |
+
zc.set_default_dtype(np.float32)
|
| 115 |
+
try:
|
| 116 |
+
np.random.seed(5)
|
| 117 |
+
model = ZeroX(vocab_size=9, n_embd=16, n_head=2, n_slot=6).eval()
|
| 118 |
+
idx = rng.integers(0, 9, size=(1, 1000))
|
| 119 |
+
par = model(idx).data
|
| 120 |
+
assert par.dtype == np.float32 and np.isfinite(par).all()
|
| 121 |
+
assert np.abs(par - model.forward_recurrent(idx)).max() < 1e-3
|
| 122 |
+
finally:
|
| 123 |
+
zc.set_default_dtype(np.float64)
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
def test_learns_to_copy_from_far_back():
|
| 127 |
+
"""记忆任务:开头一个字,中间 40 个干扰字(跨块),最后复述开头那个字。"""
|
| 128 |
+
np.random.seed(4)
|
| 129 |
+
V, gap = 6, 40
|
| 130 |
+
model = ZeroX(vocab_size=V + 1, n_embd=32, n_head=2, n_slot=8, n_layer=1)
|
| 131 |
+
opt = zc.optim.AdamW(model.parameters(), lr=0.01)
|
| 132 |
+
for _ in range(300):
|
| 133 |
+
first = np.random.randint(0, V, 32)
|
| 134 |
+
idx = np.concatenate([first[:, None], np.full((32, gap), V)], axis=1)
|
| 135 |
+
loss = zc.functional.cross_entropy(model(idx)[np.arange(32), gap], first)
|
| 136 |
+
opt.zero_grad()
|
| 137 |
+
loss.backward()
|
| 138 |
+
opt.step()
|
| 139 |
+
assert loss.item() < 0.05, loss.item()
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
if __name__ == "__main__":
|
| 143 |
+
tests = [(k, v) for k, v in list(globals().items()) if k.startswith("test_")]
|
| 144 |
+
failed = 0
|
| 145 |
+
for name, fn in tests:
|
| 146 |
+
try:
|
| 147 |
+
fn()
|
| 148 |
+
print(f" ok {name}")
|
| 149 |
+
except Exception as e: # noqa: BLE001
|
| 150 |
+
failed += 1
|
| 151 |
+
print(f" FAIL {name}: {type(e).__name__}: {e}")
|
| 152 |
+
print(f"\n{len(tests) - failed}/{len(tests)} 通过")
|
| 153 |
+
sys.exit(1 if failed else 0)
|
zerocg/__init__.py
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""zerocg: 只用 numpy 实现的迷你深度学习框架。
|
| 2 |
+
|
| 3 |
+
import zerocg as zc
|
| 4 |
+
x = zc.randn(4, 3, requires_grad=True)
|
| 5 |
+
loss = (x * x).sum()
|
| 6 |
+
loss.backward() # x.grad == 2 * x
|
| 7 |
+
"""
|
| 8 |
+
from . import functional, nn, optim
|
| 9 |
+
from .tensor import Tensor, concat, count_flops, get_default_dtype, no_grad, ones, randn, set_default_dtype, tensor, zeros
|
| 10 |
+
|
| 11 |
+
__all__ = ["Tensor", "concat", "count_flops", "no_grad", "set_default_dtype", "get_default_dtype", "tensor", "zeros", "ones", "randn", "nn", "optim", "functional"]
|
zerocg/functional.py
ADDED
|
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""常用函数。softmax / log_softmax / gelu 写成融合算子(手写反向),更快也更数值稳定。"""
|
| 2 |
+
import numpy as np
|
| 3 |
+
|
| 4 |
+
from .tensor import Tensor
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
def relu(x):
|
| 8 |
+
return x.relu()
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def tanh(x):
|
| 12 |
+
return x.tanh()
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def sigmoid_np(a):
|
| 16 |
+
"""数值稳定:只对 -|a| 求 exp,永远不会溢出。"""
|
| 17 |
+
e = np.exp(-np.abs(a))
|
| 18 |
+
return np.where(a >= 0, 1 / (1 + e), e / (1 + e))
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def sigmoid(x):
|
| 22 |
+
out = sigmoid_np(x.data)
|
| 23 |
+
return Tensor._op(out, (x,), lambda g: (g * out * (1 - out),))
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def gelu(x):
|
| 27 |
+
"""GELU(tanh 近似,GPT-2 同款)。"""
|
| 28 |
+
c = np.sqrt(2 / np.pi)
|
| 29 |
+
a = x.data
|
| 30 |
+
inner = c * (a + 0.044715 * a ** 3)
|
| 31 |
+
t = np.tanh(inner)
|
| 32 |
+
out = 0.5 * a * (1 + t)
|
| 33 |
+
|
| 34 |
+
def backward(g):
|
| 35 |
+
dinner = c * (1 + 3 * 0.044715 * a ** 2)
|
| 36 |
+
return (g * (0.5 * (1 + t) + 0.5 * a * (1 - t ** 2) * dinner),)
|
| 37 |
+
|
| 38 |
+
return Tensor._op(out.astype(a.dtype), (x,), backward)
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def softmax(x, axis=-1):
|
| 42 |
+
e = np.exp(x.data - x.data.max(axis=axis, keepdims=True))
|
| 43 |
+
s = e / e.sum(axis=axis, keepdims=True)
|
| 44 |
+
return Tensor._op(s, (x,), lambda g: (s * (g - (g * s).sum(axis=axis, keepdims=True)),))
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def log_softmax(x, axis=-1):
|
| 48 |
+
shifted = x.data - x.data.max(axis=axis, keepdims=True)
|
| 49 |
+
out = shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))
|
| 50 |
+
s = np.exp(out)
|
| 51 |
+
return Tensor._op(out, (x,), lambda g: (g - s * g.sum(axis=axis, keepdims=True),))
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def cross_entropy(logits, targets):
|
| 55 |
+
"""logits: (N, C);targets: (N,) 整数类别。返回平均负对数似然。"""
|
| 56 |
+
targets = targets.data if isinstance(targets, Tensor) else np.asarray(targets)
|
| 57 |
+
logp = log_softmax(logits, axis=-1)
|
| 58 |
+
return -logp[np.arange(len(targets)), targets].mean()
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def mse_loss(pred, target):
|
| 62 |
+
return ((pred - target) ** 2).mean()
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def dropout(x, p, training=True):
|
| 66 |
+
if not training or p == 0:
|
| 67 |
+
return x
|
| 68 |
+
mask = (np.random.rand(*x.shape) >= p).astype(x.dtype) / (1 - p)
|
| 69 |
+
return x * Tensor(mask)
|
zerocg/nn.py
ADDED
|
@@ -0,0 +1,162 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""神经网络层。Module 会自动收集属性里的 Parameter / 子 Module。"""
|
| 2 |
+
import numpy as np
|
| 3 |
+
|
| 4 |
+
from . import functional as F
|
| 5 |
+
from .tensor import Tensor, get_default_dtype
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
class Parameter(Tensor):
|
| 9 |
+
def __init__(self, data):
|
| 10 |
+
super().__init__(np.asarray(data, dtype=get_default_dtype()), requires_grad=True)
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
class Module:
|
| 14 |
+
training = True
|
| 15 |
+
|
| 16 |
+
def forward(self, *args, **kwargs):
|
| 17 |
+
raise NotImplementedError
|
| 18 |
+
|
| 19 |
+
def __call__(self, *args, **kwargs):
|
| 20 |
+
return self.forward(*args, **kwargs)
|
| 21 |
+
|
| 22 |
+
def named_parameters(self, prefix=""):
|
| 23 |
+
for name, v in vars(self).items():
|
| 24 |
+
items = [(name, v)] if not isinstance(v, (list, tuple)) else [(f"{name}.{i}", x) for i, x in enumerate(v)]
|
| 25 |
+
for n, x in items:
|
| 26 |
+
if isinstance(x, Parameter):
|
| 27 |
+
yield prefix + n, x
|
| 28 |
+
elif isinstance(x, Module):
|
| 29 |
+
yield from x.named_parameters(prefix + n + ".")
|
| 30 |
+
|
| 31 |
+
def parameters(self):
|
| 32 |
+
return [p for _, p in self.named_parameters()]
|
| 33 |
+
|
| 34 |
+
def modules(self):
|
| 35 |
+
yield self
|
| 36 |
+
for v in vars(self).values():
|
| 37 |
+
for x in (v if isinstance(v, (list, tuple)) else [v]):
|
| 38 |
+
if isinstance(x, Module):
|
| 39 |
+
yield from x.modules()
|
| 40 |
+
|
| 41 |
+
def train(self, mode=True):
|
| 42 |
+
for m in self.modules():
|
| 43 |
+
m.training = mode
|
| 44 |
+
return self
|
| 45 |
+
|
| 46 |
+
def eval(self):
|
| 47 |
+
return self.train(False)
|
| 48 |
+
|
| 49 |
+
def zero_grad(self):
|
| 50 |
+
for p in self.parameters():
|
| 51 |
+
p.grad = None
|
| 52 |
+
|
| 53 |
+
def num_parameters(self):
|
| 54 |
+
return sum(p.data.size for p in self.parameters())
|
| 55 |
+
|
| 56 |
+
# ---------- 保存 / 加载 ----------
|
| 57 |
+
def state_dict(self):
|
| 58 |
+
return {n: p.data.copy() for n, p in self.named_parameters()}
|
| 59 |
+
|
| 60 |
+
def load_state_dict(self, state):
|
| 61 |
+
for n, p in self.named_parameters():
|
| 62 |
+
assert state[n].shape == p.shape, f"{n} 形状不匹配: {state[n].shape} vs {p.shape}"
|
| 63 |
+
p.data = state[n].astype(p.dtype)
|
| 64 |
+
|
| 65 |
+
def save(self, path):
|
| 66 |
+
np.savez(path, **self.state_dict())
|
| 67 |
+
|
| 68 |
+
def load(self, path):
|
| 69 |
+
with np.load(path) as f:
|
| 70 |
+
self.load_state_dict(dict(f))
|
| 71 |
+
return self
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
class Linear(Module):
|
| 75 |
+
def __init__(self, in_features, out_features, bias=True):
|
| 76 |
+
bound = 1 / np.sqrt(in_features)
|
| 77 |
+
self.weight = Parameter(np.random.uniform(-bound, bound, (in_features, out_features)))
|
| 78 |
+
self.bias = Parameter(np.zeros(out_features)) if bias else None
|
| 79 |
+
|
| 80 |
+
def forward(self, x):
|
| 81 |
+
y = x @ self.weight
|
| 82 |
+
return y + self.bias if self.bias is not None else y
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
class Embedding(Module):
|
| 86 |
+
def __init__(self, num_embeddings, dim):
|
| 87 |
+
self.weight = Parameter(np.random.randn(num_embeddings, dim) * 0.02)
|
| 88 |
+
|
| 89 |
+
def forward(self, idx):
|
| 90 |
+
return self.weight[idx.data if isinstance(idx, Tensor) else np.asarray(idx)]
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
class LayerNorm(Module):
|
| 94 |
+
def __init__(self, dim, eps=1e-5):
|
| 95 |
+
self.eps = eps
|
| 96 |
+
self.weight = Parameter(np.ones(dim))
|
| 97 |
+
self.bias = Parameter(np.zeros(dim))
|
| 98 |
+
|
| 99 |
+
def forward(self, x):
|
| 100 |
+
mu = x.mean(axis=-1, keepdims=True)
|
| 101 |
+
xc = x - mu
|
| 102 |
+
var = (xc * xc).mean(axis=-1, keepdims=True)
|
| 103 |
+
return xc / (var + self.eps).sqrt() * self.weight + self.bias
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
class Dropout(Module):
|
| 107 |
+
def __init__(self, p=0.1):
|
| 108 |
+
self.p = p
|
| 109 |
+
|
| 110 |
+
def forward(self, x):
|
| 111 |
+
return F.dropout(x, self.p, self.training)
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class ReLU(Module):
|
| 115 |
+
def forward(self, x):
|
| 116 |
+
return x.relu()
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
class Tanh(Module):
|
| 120 |
+
def forward(self, x):
|
| 121 |
+
return x.tanh()
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
class GELU(Module):
|
| 125 |
+
def forward(self, x):
|
| 126 |
+
return F.gelu(x)
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
class Sequential(Module):
|
| 130 |
+
def __init__(self, *layers):
|
| 131 |
+
self.layers = list(layers)
|
| 132 |
+
|
| 133 |
+
def forward(self, x):
|
| 134 |
+
for layer in self.layers:
|
| 135 |
+
x = layer(x)
|
| 136 |
+
return x
|
| 137 |
+
|
| 138 |
+
def __getitem__(self, i):
|
| 139 |
+
return self.layers[i]
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
class CausalSelfAttention(Module):
|
| 143 |
+
"""多头因果自注意力:位置 t 只能看到 <= t 的位置。输入/输出形状 (B, T, C)。"""
|
| 144 |
+
|
| 145 |
+
def __init__(self, n_embd, n_head, dropout=0.0):
|
| 146 |
+
assert n_embd % n_head == 0
|
| 147 |
+
self.n_head = n_head
|
| 148 |
+
self.qkv = Linear(n_embd, 3 * n_embd)
|
| 149 |
+
self.proj = Linear(n_embd, n_embd)
|
| 150 |
+
self.attn_drop = Dropout(dropout)
|
| 151 |
+
self.resid_drop = Dropout(dropout)
|
| 152 |
+
|
| 153 |
+
def forward(self, x):
|
| 154 |
+
B, T, C = x.shape
|
| 155 |
+
H, D = self.n_head, C // self.n_head
|
| 156 |
+
qkv = self.qkv(x).reshape(B, T, 3, H, D).transpose(2, 0, 3, 1, 4) # (3, B, H, T, D)
|
| 157 |
+
q, k, v = qkv[0], qkv[1], qkv[2]
|
| 158 |
+
att = (q @ k.transpose(0, 1, 3, 2)) * (1.0 / np.sqrt(D)) # (B, H, T, T)
|
| 159 |
+
att = att.masked_fill(np.triu(np.ones((T, T), dtype=bool), 1), -1e9)
|
| 160 |
+
att = self.attn_drop(F.softmax(att, axis=-1))
|
| 161 |
+
y = (att @ v).transpose(0, 2, 1, 3).reshape(B, T, C)
|
| 162 |
+
return self.resid_drop(self.proj(y))
|
zerocg/optim.py
ADDED
|
@@ -0,0 +1,87 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""优化器。"""
|
| 2 |
+
import numpy as np
|
| 3 |
+
|
| 4 |
+
|
| 5 |
+
class Optimizer:
|
| 6 |
+
def __init__(self, params, lr):
|
| 7 |
+
self.params = list(params)
|
| 8 |
+
self.lr = lr
|
| 9 |
+
|
| 10 |
+
def zero_grad(self):
|
| 11 |
+
for p in self.params:
|
| 12 |
+
p.grad = None
|
| 13 |
+
|
| 14 |
+
def step(self):
|
| 15 |
+
raise NotImplementedError
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class SGD(Optimizer):
|
| 19 |
+
def __init__(self, params, lr=0.01, momentum=0.0):
|
| 20 |
+
super().__init__(params, lr)
|
| 21 |
+
self.momentum = momentum
|
| 22 |
+
self.v = [np.zeros_like(p.data) for p in self.params]
|
| 23 |
+
|
| 24 |
+
def step(self):
|
| 25 |
+
for p, v in zip(self.params, self.v):
|
| 26 |
+
if p.grad is None:
|
| 27 |
+
continue
|
| 28 |
+
v *= self.momentum
|
| 29 |
+
v += p.grad
|
| 30 |
+
p.data -= self.lr * v
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
class AdamW(Optimizer):
|
| 34 |
+
"""Adam + 解耦权重衰减。weight_decay=0 时就是普通 Adam。"""
|
| 35 |
+
|
| 36 |
+
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
|
| 37 |
+
super().__init__(params, lr)
|
| 38 |
+
self.b1, self.b2 = betas
|
| 39 |
+
self.eps = eps
|
| 40 |
+
self.weight_decay = weight_decay
|
| 41 |
+
self.t = 0
|
| 42 |
+
self.m = [np.zeros_like(p.data) for p in self.params]
|
| 43 |
+
self.v = [np.zeros_like(p.data) for p in self.params]
|
| 44 |
+
|
| 45 |
+
def step(self):
|
| 46 |
+
self.t += 1
|
| 47 |
+
c1, c2 = 1 - self.b1 ** self.t, 1 - self.b2 ** self.t
|
| 48 |
+
for p, m, v in zip(self.params, self.m, self.v):
|
| 49 |
+
if p.grad is None:
|
| 50 |
+
continue
|
| 51 |
+
g = p.grad
|
| 52 |
+
m *= self.b1
|
| 53 |
+
m += (1 - self.b1) * g
|
| 54 |
+
v *= self.b2
|
| 55 |
+
v += (1 - self.b2) * g * g
|
| 56 |
+
if self.weight_decay:
|
| 57 |
+
p.data -= self.lr * self.weight_decay * p.data
|
| 58 |
+
p.data -= self.lr * (m / c1) / (np.sqrt(v / c2) + self.eps)
|
| 59 |
+
|
| 60 |
+
def state_dict(self):
|
| 61 |
+
"""优化器状态(步数 + 一阶/二阶动量)。和 model.state_dict() 一起存,才能精确续训。"""
|
| 62 |
+
state = {"t": np.array(self.t), "lr": np.array(self.lr)}
|
| 63 |
+
for i, (m, v) in enumerate(zip(self.m, self.v)):
|
| 64 |
+
state[f"m.{i}"], state[f"v.{i}"] = m.copy(), v.copy()
|
| 65 |
+
return state
|
| 66 |
+
|
| 67 |
+
def load_state_dict(self, state):
|
| 68 |
+
assert sum(k.startswith("m.") for k in state) == len(self.params), "参数个数和存档对不上"
|
| 69 |
+
self.t, self.lr = int(state["t"]), float(state["lr"])
|
| 70 |
+
for i in range(len(self.params)):
|
| 71 |
+
self.m[i][...] = state[f"m.{i}"]
|
| 72 |
+
self.v[i][...] = state[f"v.{i}"]
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
Adam = AdamW
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def clip_grad_norm(params, max_norm):
|
| 79 |
+
"""按全局 L2 范数裁剪梯度,返回裁剪前的范数。"""
|
| 80 |
+
grads = [p.grad for p in params if p.grad is not None]
|
| 81 |
+
total = float(np.sqrt(sum((g.astype(np.float64) ** 2).sum() for g in grads)))
|
| 82 |
+
if total > max_norm:
|
| 83 |
+
scale = max_norm / (total + 1e-6)
|
| 84 |
+
for p in params:
|
| 85 |
+
if p.grad is not None:
|
| 86 |
+
p.grad = p.grad * scale
|
| 87 |
+
return total
|
zerocg/tensor.py
ADDED
|
@@ -0,0 +1,320 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tensor + 反向模式自动求导(只依赖 numpy)。
|
| 2 |
+
|
| 3 |
+
每个运算产出一个新 Tensor,并记录:
|
| 4 |
+
_parents : 参与运算的输入 Tensor
|
| 5 |
+
_backward : 函数 g -> (对每个 parent 的梯度, ...)
|
| 6 |
+
backward() 先拓扑排序,再从输出往回逐个调用 _backward,把梯度累加到 .grad。
|
| 7 |
+
"""
|
| 8 |
+
import numpy as np
|
| 9 |
+
|
| 10 |
+
_grad_enabled = True
|
| 11 |
+
_default_dtype = np.dtype(np.float32)
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def set_default_dtype(dtype):
|
| 15 |
+
"""所有浮点数据进入框架时统一转成这个类型(默认 float32)。
|
| 16 |
+
只有测试里的数值梯度检查需要切到 float64。整数数组(下标)不受影响。"""
|
| 17 |
+
global _default_dtype
|
| 18 |
+
_default_dtype = np.dtype(dtype)
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def get_default_dtype():
|
| 22 |
+
return _default_dtype
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
_flops = None # count_flops() 打开时是一个 [计数],否则 None
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
class count_flops:
|
| 29 |
+
"""with count_flops() as c: ... 之后 c.flops = 期间所有矩阵乘的浮点运算次数(乘加算 2 次)。
|
| 30 |
+
只统计矩阵乘(模型里绝大部分计算量),不统计逐元素运算;不受机器负载影响,比计时可靠。"""
|
| 31 |
+
|
| 32 |
+
def __enter__(self):
|
| 33 |
+
global _flops
|
| 34 |
+
self._prev, _flops = _flops, [0]
|
| 35 |
+
return self
|
| 36 |
+
|
| 37 |
+
def __exit__(self, *exc):
|
| 38 |
+
global _flops
|
| 39 |
+
self.flops = _flops[0]
|
| 40 |
+
_flops = self._prev
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
class no_grad:
|
| 44 |
+
"""with no_grad(): 内部的运算不建计算图(推理/评估时用,省内存)。"""
|
| 45 |
+
|
| 46 |
+
def __enter__(self):
|
| 47 |
+
global _grad_enabled
|
| 48 |
+
self._prev, _grad_enabled = _grad_enabled, False
|
| 49 |
+
|
| 50 |
+
def __exit__(self, *exc):
|
| 51 |
+
global _grad_enabled
|
| 52 |
+
_grad_enabled = self._prev
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def _unbroadcast(grad, shape):
|
| 56 |
+
"""把广播后的梯度求和还原成原始形状。"""
|
| 57 |
+
while grad.ndim > len(shape):
|
| 58 |
+
grad = grad.sum(axis=0)
|
| 59 |
+
for i, s in enumerate(shape):
|
| 60 |
+
if s == 1 and grad.shape[i] != 1:
|
| 61 |
+
grad = grad.sum(axis=i, keepdims=True)
|
| 62 |
+
return grad
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def _to_array(x):
|
| 66 |
+
if isinstance(x, Tensor):
|
| 67 |
+
return x.data
|
| 68 |
+
arr = np.asarray(x)
|
| 69 |
+
# 浮点一律转成默认类型,避免 float32 和 float64 混在一起被 numpy 悄悄升成 float64
|
| 70 |
+
if arr.dtype.kind == "f" and arr.dtype != _default_dtype:
|
| 71 |
+
arr = arr.astype(_default_dtype)
|
| 72 |
+
return arr
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
class Tensor:
|
| 76 |
+
__array_priority__ = 100 # 让 ndarray + Tensor 走 Tensor.__radd__
|
| 77 |
+
|
| 78 |
+
def __init__(self, data, requires_grad=False):
|
| 79 |
+
self.data = _to_array(data)
|
| 80 |
+
self.requires_grad = requires_grad
|
| 81 |
+
self.grad = None
|
| 82 |
+
self._parents = ()
|
| 83 |
+
self._backward = None
|
| 84 |
+
|
| 85 |
+
# ---------- 基础属性 ----------
|
| 86 |
+
@property
|
| 87 |
+
def shape(self):
|
| 88 |
+
return self.data.shape
|
| 89 |
+
|
| 90 |
+
@property
|
| 91 |
+
def ndim(self):
|
| 92 |
+
return self.data.ndim
|
| 93 |
+
|
| 94 |
+
@property
|
| 95 |
+
def dtype(self):
|
| 96 |
+
return self.data.dtype
|
| 97 |
+
|
| 98 |
+
def __len__(self):
|
| 99 |
+
return len(self.data)
|
| 100 |
+
|
| 101 |
+
def __repr__(self):
|
| 102 |
+
return f"Tensor({self.data}, requires_grad={self.requires_grad})"
|
| 103 |
+
|
| 104 |
+
def numpy(self):
|
| 105 |
+
return self.data
|
| 106 |
+
|
| 107 |
+
def item(self):
|
| 108 |
+
return self.data.item()
|
| 109 |
+
|
| 110 |
+
def detach(self):
|
| 111 |
+
return Tensor(self.data)
|
| 112 |
+
|
| 113 |
+
# ---------- 建图 ----------
|
| 114 |
+
@staticmethod
|
| 115 |
+
def _op(data, parents, backward):
|
| 116 |
+
"""创建运算结果。backward(g) 返回与 parents 一一对应的梯度(可为 None)。"""
|
| 117 |
+
out = Tensor(data)
|
| 118 |
+
if _grad_enabled and any(p.requires_grad for p in parents):
|
| 119 |
+
out.requires_grad = True
|
| 120 |
+
out._parents = parents
|
| 121 |
+
out._backward = backward
|
| 122 |
+
return out
|
| 123 |
+
|
| 124 |
+
def backward(self, grad=None):
|
| 125 |
+
if grad is None:
|
| 126 |
+
if self.data.size != 1:
|
| 127 |
+
raise RuntimeError("非标量输出调用 backward() 需要显式传入 grad")
|
| 128 |
+
grad = np.ones_like(self.data)
|
| 129 |
+
|
| 130 |
+
# 迭代式拓扑排序(避免深图递归爆栈)
|
| 131 |
+
topo, visited, stack = [], set(), [(self, False)]
|
| 132 |
+
while stack:
|
| 133 |
+
node, expanded = stack.pop()
|
| 134 |
+
if expanded:
|
| 135 |
+
topo.append(node)
|
| 136 |
+
continue
|
| 137 |
+
if id(node) in visited:
|
| 138 |
+
continue
|
| 139 |
+
visited.add(id(node))
|
| 140 |
+
stack.append((node, True))
|
| 141 |
+
for p in node._parents:
|
| 142 |
+
if p.requires_grad and id(p) not in visited:
|
| 143 |
+
stack.append((p, False))
|
| 144 |
+
|
| 145 |
+
self.grad = grad if self.grad is None else self.grad + grad
|
| 146 |
+
for node in reversed(topo):
|
| 147 |
+
if node._backward is None or node.grad is None:
|
| 148 |
+
continue
|
| 149 |
+
for p, g in zip(node._parents, node._backward(node.grad)):
|
| 150 |
+
if g is not None and p.requires_grad:
|
| 151 |
+
p.grad = g if p.grad is None else p.grad + g
|
| 152 |
+
# 释放计算图,中间节点的梯度也不再需要
|
| 153 |
+
node._parents, node._backward, node.grad = (), None, None
|
| 154 |
+
|
| 155 |
+
# ---------- 逐元素运算 ----------
|
| 156 |
+
def __add__(self, other):
|
| 157 |
+
other = other if isinstance(other, Tensor) else Tensor(other)
|
| 158 |
+
a, b = self, other
|
| 159 |
+
return Tensor._op(a.data + b.data, (a, b),
|
| 160 |
+
lambda g: (_unbroadcast(g, a.shape), _unbroadcast(g, b.shape)))
|
| 161 |
+
|
| 162 |
+
def __mul__(self, other):
|
| 163 |
+
other = other if isinstance(other, Tensor) else Tensor(other)
|
| 164 |
+
a, b = self, other
|
| 165 |
+
return Tensor._op(a.data * b.data, (a, b),
|
| 166 |
+
lambda g: (_unbroadcast(g * b.data, a.shape), _unbroadcast(g * a.data, b.shape)))
|
| 167 |
+
|
| 168 |
+
def __truediv__(self, other):
|
| 169 |
+
other = other if isinstance(other, Tensor) else Tensor(other)
|
| 170 |
+
a, b = self, other
|
| 171 |
+
return Tensor._op(a.data / b.data, (a, b),
|
| 172 |
+
lambda g: (_unbroadcast(g / b.data, a.shape),
|
| 173 |
+
_unbroadcast(-g * a.data / (b.data ** 2), b.shape)))
|
| 174 |
+
|
| 175 |
+
def __pow__(self, p):
|
| 176 |
+
assert isinstance(p, (int, float)), "只支持标量指数"
|
| 177 |
+
a = self
|
| 178 |
+
return Tensor._op(a.data ** p, (a,), lambda g: (g * p * a.data ** (p - 1),))
|
| 179 |
+
|
| 180 |
+
def __neg__(self):
|
| 181 |
+
return Tensor._op(-self.data, (self,), lambda g: (-g,))
|
| 182 |
+
|
| 183 |
+
def __sub__(self, other):
|
| 184 |
+
return self + (-other if isinstance(other, Tensor) else Tensor(-_to_array(other)))
|
| 185 |
+
|
| 186 |
+
def __radd__(self, other):
|
| 187 |
+
return self + other
|
| 188 |
+
|
| 189 |
+
def __rmul__(self, other):
|
| 190 |
+
return self * other
|
| 191 |
+
|
| 192 |
+
def __rsub__(self, other):
|
| 193 |
+
return (-self) + other
|
| 194 |
+
|
| 195 |
+
def __rtruediv__(self, other):
|
| 196 |
+
return Tensor(other) / self
|
| 197 |
+
|
| 198 |
+
def exp(self):
|
| 199 |
+
out = np.exp(self.data)
|
| 200 |
+
return Tensor._op(out, (self,), lambda g: (g * out,))
|
| 201 |
+
|
| 202 |
+
def log(self):
|
| 203 |
+
a = self
|
| 204 |
+
return Tensor._op(np.log(a.data), (a,), lambda g: (g / a.data,))
|
| 205 |
+
|
| 206 |
+
def sqrt(self):
|
| 207 |
+
return self ** 0.5
|
| 208 |
+
|
| 209 |
+
def tanh(self):
|
| 210 |
+
out = np.tanh(self.data)
|
| 211 |
+
return Tensor._op(out, (self,), lambda g: (g * (1 - out ** 2),))
|
| 212 |
+
|
| 213 |
+
def relu(self):
|
| 214 |
+
mask = self.data > 0
|
| 215 |
+
return Tensor._op(self.data * mask, (self,), lambda g: (g * mask,))
|
| 216 |
+
|
| 217 |
+
def masked_fill(self, mask, value):
|
| 218 |
+
mask = _to_array(mask).astype(bool)
|
| 219 |
+
return Tensor._op(np.where(mask, value, self.data).astype(self.dtype), (self,),
|
| 220 |
+
lambda g: (np.where(mask, 0, g),))
|
| 221 |
+
|
| 222 |
+
# ---------- 归约 ----------
|
| 223 |
+
def sum(self, axis=None, keepdims=False):
|
| 224 |
+
a = self
|
| 225 |
+
|
| 226 |
+
def backward(g):
|
| 227 |
+
if axis is not None and not keepdims:
|
| 228 |
+
g = np.expand_dims(g, axis)
|
| 229 |
+
return (np.broadcast_to(g, a.shape).copy(),)
|
| 230 |
+
|
| 231 |
+
return Tensor._op(a.data.sum(axis=axis, keepdims=keepdims), (a,), backward)
|
| 232 |
+
|
| 233 |
+
def mean(self, axis=None, keepdims=False):
|
| 234 |
+
n = self.data.size if axis is None else np.prod([self.shape[i] for i in np.atleast_1d(axis)])
|
| 235 |
+
return self.sum(axis, keepdims) * (1.0 / n)
|
| 236 |
+
|
| 237 |
+
def max(self, axis=None, keepdims=False):
|
| 238 |
+
a = self
|
| 239 |
+
out = a.data.max(axis=axis, keepdims=True)
|
| 240 |
+
|
| 241 |
+
def backward(g):
|
| 242 |
+
if axis is not None and not keepdims:
|
| 243 |
+
g = np.expand_dims(g, axis)
|
| 244 |
+
mask = a.data == out
|
| 245 |
+
return (mask * g / mask.sum(axis=axis, keepdims=True),) # 并列最大值平分梯度
|
| 246 |
+
|
| 247 |
+
return Tensor._op(out if keepdims else out.squeeze(axis), (a,), backward)
|
| 248 |
+
|
| 249 |
+
# ---------- 形状 ----------
|
| 250 |
+
def reshape(self, *shape):
|
| 251 |
+
a = self
|
| 252 |
+
return Tensor._op(a.data.reshape(*shape), (a,), lambda g: (g.reshape(a.shape),))
|
| 253 |
+
|
| 254 |
+
def transpose(self, *axes):
|
| 255 |
+
axes = axes or tuple(reversed(range(self.ndim)))
|
| 256 |
+
inv = np.argsort(axes)
|
| 257 |
+
return Tensor._op(self.data.transpose(axes), (self,), lambda g: (g.transpose(inv),))
|
| 258 |
+
|
| 259 |
+
@property
|
| 260 |
+
def T(self):
|
| 261 |
+
return self.transpose()
|
| 262 |
+
|
| 263 |
+
def __getitem__(self, idx):
|
| 264 |
+
a = self
|
| 265 |
+
if isinstance(idx, Tensor):
|
| 266 |
+
idx = idx.data
|
| 267 |
+
|
| 268 |
+
parts = idx if isinstance(idx, tuple) else (idx,)
|
| 269 |
+
fancy = any(isinstance(i, (np.ndarray, list)) for i in parts)
|
| 270 |
+
|
| 271 |
+
def backward(g):
|
| 272 |
+
grad = np.zeros_like(a.data)
|
| 273 |
+
if fancy:
|
| 274 |
+
np.add.at(grad, idx, g) # 数组下标可能重复(Embedding 查表),必须累加
|
| 275 |
+
else:
|
| 276 |
+
grad[idx] = g # 普通切片不会重复,直接写,比 add.at 快得多
|
| 277 |
+
return (grad,)
|
| 278 |
+
|
| 279 |
+
return Tensor._op(a.data[idx], (a,), backward)
|
| 280 |
+
|
| 281 |
+
# ---------- 矩阵乘 ----------
|
| 282 |
+
def __matmul__(self, other):
|
| 283 |
+
other = other if isinstance(other, Tensor) else Tensor(other)
|
| 284 |
+
a, b = self, other
|
| 285 |
+
assert a.ndim >= 2 and b.ndim >= 2, "matmul 需要至少 2 维"
|
| 286 |
+
out = a.data @ b.data
|
| 287 |
+
if _flops is not None:
|
| 288 |
+
_flops[0] += 2 * out.size * a.shape[-1]
|
| 289 |
+
|
| 290 |
+
def backward(g):
|
| 291 |
+
ga = g @ np.swapaxes(b.data, -1, -2)
|
| 292 |
+
gb = np.swapaxes(a.data, -1, -2) @ g
|
| 293 |
+
return _unbroadcast(ga, a.shape), _unbroadcast(gb, b.shape)
|
| 294 |
+
|
| 295 |
+
return Tensor._op(out, (a, b), backward)
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
def concat(tensors, axis=0):
|
| 299 |
+
"""沿 axis 拼接。反向时把梯度按原来的大小切回去。"""
|
| 300 |
+
tensors = tuple(t if isinstance(t, Tensor) else Tensor(t) for t in tensors)
|
| 301 |
+
splits = np.cumsum([t.shape[axis] for t in tensors])[:-1]
|
| 302 |
+
return Tensor._op(np.concatenate([t.data for t in tensors], axis=axis), tensors,
|
| 303 |
+
lambda g: tuple(np.split(g, splits, axis=axis)))
|
| 304 |
+
|
| 305 |
+
|
| 306 |
+
# ---------- 构造函数 ----------
|
| 307 |
+
def tensor(data, requires_grad=False):
|
| 308 |
+
return Tensor(data, requires_grad)
|
| 309 |
+
|
| 310 |
+
|
| 311 |
+
def zeros(*shape, requires_grad=False):
|
| 312 |
+
return Tensor(np.zeros(shape, dtype=_default_dtype), requires_grad)
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
def ones(*shape, requires_grad=False):
|
| 316 |
+
return Tensor(np.ones(shape, dtype=_default_dtype), requires_grad)
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
def randn(*shape, requires_grad=False):
|
| 320 |
+
return Tensor(np.random.randn(*shape).astype(_default_dtype), requires_grad)
|