Squash commits

Author: jerrywliu

.gitignore

@@ -0,0 +1,18 @@
+.ipynb_checkpoints
+__pycache__
+*/__pycache__/
+models.zip
+/models/
+/data/
+/figs/
+
+.DS_Store
+src/slurm*.out
+src/figs/
+src/*.png
+notebooks/figs/
+notebooks/**/figs/
+notebooks/**/*.png
+src/datagen/test/figs/
+src/conf/wandb.yaml
+scripts/*.png

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 HazyResearch
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,63 @@
+# Towards Learning High-Precision Least Squares Algorithms with Sequence Models
+
+![Prior work focuses on statistical least squares: Transformers approximate Bayes-optimal estimators (left, adapted from Garg et al. 2022). In this work, we focus on numerical least squares: Transformers struggle to obtain precise solutions (inset). Using a high-precision training recipe, we train two polynomial architectures to perform high-precision gradient descent iterates on least squares (right): applied iteratively, they reach $10^{-13}$ MSE.](assets/banner_fig.png)
+
+This repository contains code for the following paper:
+
+> **Towards Learning High-Precision Least Squares Algorithms with Sequence Models.**
+>
+> Jerry Liu, Jessica Grogan, Owen Dugan, Ashish Rao, Simran Arora, Atri Rudra, Chris Ré.
+> ICLR 2025.
+
+## Dependencies
+Install dependencies with
+```
+conda create -n "precision-ls" python=3.10
+conda activate precision-ls
+pip install -r requirements.txt
+```
+
+## Code structure
+The code is organized as follows:
+- `notebooks/`: contains notebooks for the experiments
+- `replication_scripts/`: contains scripts for running the experiments
+- `src/datagen/`: contains code for data generation (tasks and samplers)
+- `src/models/`: contains code for the models, including different sequence mixers
+- `src/schedulers/`: contains code for the learning rate schedulers, including the adaptive one used in the paper
+
+## Running experiments
+To test the data generation, run
+```
+bash src/datagen/test.sh
+```
+
+To run the least squares experiments, run
+```
+bash replication_scripts/least_squares.sh
+```
+
+To run the experiments and ablations with linear algebra primitives, run
+```
+bash replication_scripts/primitives.sh
+```
+
+To run the experiments and ablations with the explicit gradient and k-iterate gradient descent tasks, run
+```
+bash replication_scripts/explicit_gradient.sh
+```
+
+To run the experiments with in-context ODEs, run
+```
+bash replication_scripts/odes.sh
+```
+
+## Citation
+If you find this work useful, please cite it as follows:
+```
+@article{TODO,
+  title={TODO},
+  author={TODO},
+  journal={TODO},
+  year={TODO},
+}
+```