JG1310/repro-chain-of-thought-gradient-descent-code
Reproduction bundle — Chain-of-Thought Gradient Descent (ICML 2026, submission #443) OpenReview: https://openreview.net/forum?id=uZ8JZ1Lw9a An independent reproduction. No official code, checkpoints, or data were released with the paper, so everything here is implemented from the paper text (Definitions 2.2–2.3, 4.3–4.4; Theorems 4.1–4.2; Propositions 4.1–4.4; Appendix D and E). What is here scripts/ common.py N-layer FFN forward/backprop (Def… See the full description on the dataset page: https://huggingface.co/datasets/JG1310/repro-chain-of-thought-gradient-descent-code.
Reproduction bundle — Chain-of-Thought Gradient Descent (ICML 2026, submission #443)
OpenReview: <https://openreview.net/forum?id=uZ8JZ1Lw9a>
An independent reproduction. No official code, checkpoints, or data were released with the paper, so everything here is implemented from the paper text (Definitions 2.2–2.3, 4.3–4.4; Theorems 4.1–4.2; Propositions 4.1–4.4; Appendix D and E).
What is here
scripts/
common.py N-layer FFN forward/backprop (Def 2.2/2.3), CoT block
formats F_i, B_i, X_l, Y_l (Def 4.3), dynamic-masking
rule M(t) (Def 4.4)
exp1_statistical.py Appendix E.1 — ICGD of GLM / ridge / lasso
exp2_operators.py Appendix E.2 — separate forward + backward ICGD operators
exp3_endtoend.py Appendix E.3 — end-to-end ICGD of a 3-layer FFN, and the
10-round rollout of Figure 5
exp4_masking_and_cost.py Claim 2 — masking-rule correctness vs. the Appendix D
table, block-reconstruction exactness, and the O(N) cost
benchmark against the Wu et al. (2025) packing
run_all.py driver used for the Hugging Face GPU Job
make_figures.py Plotly figures + raw CSVs from the result JSONs
make_poster_embed.py builds the logbook `poster_embed.html` from a rendered poster
outputs/ result JSONs from the GPU Job (+ run_manifest.json)
figures/ Plotly HTML figures and their raw CSVs
poster/ posterly poster source, rendered PDF/PNG, gate reportRerun
pip install torch numpy plotly
cd scripts
# Claim 2 — structural checks + cost sweep (seconds, CPU is fine)
python exp4_masking_and_cost.py --outdir ../outputs
# Claim 1 — the three Appendix E experiments (GPU recommended)
python exp1_statistical.py --outdir ../outputs --seeds 10 --steps 4000
python exp2_operators.py --outdir ../outputs --K 400 --epochs 20 --rounds 5
python exp3_endtoend.py --outdir ../outputs --seeds 10 --steps 4000
# everything at once, exactly as the GPU Job ran it
python run_all.py --outdir ../outputs # add --quick for a smoke test
python make_figures.py --indir ../outputs --outdir ../figuresOn a CPU box, export TORCH_THREADS=4 first — these are very small tensors and PyTorch's default thread count causes severe oversubscription (measured ~16x slowdown at 16 threads).
Choices the paper leaves unspecified
The paper's Appendix E omits several settings. These are ours, and are held fixed across all runs:
The last row matters. Every block in Definition 4.3 has the form (vector) * 1_{1xd}, so all d of its columns are identical and it carries R = 1 + 3d + d^2 numbers rather than R * d. We therefore represent a block as a single token of width R, which is also what Figure 1 depicts ("only 2 tokens processed").
Definition 4.3 also writes X_l with a 0_{(3d+d^2) x d} filler and Y_l with a 0_{(2d+d^2) x d} filler, which give inconsistent row counts (4d+d^2+1 vs 3d+d^2+1). We pad both to the common R and use a distinct scalar tag in the final row to mark the four routed input types that Proposition 4.3 requires.
Hub artifacts
- Code: <https://huggingface.co/datasets/JG1310/repro-chain-of-thought-gradient-descent-code>
- Logbook Space:
JG1310/repro-chain-of-thought-gradient-descent