{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/test-time-training-on-nearest-neighbors-for","title":"Test-Time Training on Nearest Neighbors for Large Language Models","arxiv_id":"2305.18466","date":"2023-05-29","proceeding":null,"authors":["Moritz Hardt","Yu Sun"],"abstract":"Many recent efforts augment language models with retrieval, by adding retrieved data to the input context. For this approach to succeed, the retrieved data must be added at both training and test time. Moreover, as input length grows linearly with the size of retrieved data, cost in computation and memory grows quadratically for modern Transformers. To avoid these complications, we simply fine-tune the model on retrieved data at test time, using its standard training setup. We build a large-scale distributed index based on text embeddings of the Pile dataset. For each test input, our system retrieves its neighbors and fine-tunes the model on their text. Surprisingly, retrieving and training on as few as 20 neighbors, each for only one gradient iteration, drastically improves performance across more than 20 language modeling tasks in the Pile. For example, test-time training with nearest neighbors significantly narrows the performance gap between a small GPT-2 and a GPT-Neo model more than 10 times larger. Sufficient index quality and size, however, are necessary. Our work establishes a first baseline of test-time training for language modeling.","url_abs":"https://arxiv.org/abs/2305.18466v3","url_pdf":"https://arxiv.org/pdf/2305.18466v3.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"test-time-training-on-nearest-neighbors-for","repo_url":"https://github.com/socialfoundations/tttlm","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"none","reach":{"status":"ok","spdx":"MIT"}}],"tasks":[{"task_slug":"language-modeling","task_name":"Language Modeling"},{"task_slug":"language-modelling","task_name":"Language Modelling"},{"task_slug":"prompt-engineering","task_name":"Prompt Engineering"},{"task_slug":"retrieval","task_name":"Retrieval"}],"methods":[{"method_slug":"adam","method_name":"Adam"},{"method_slug":"attention","method_name":"Attention"},{"method_slug":"attention-dropout","method_name":"Attention Dropout"},{"method_slug":"bpe","method_name":"BPE"},{"method_slug":"cosine-annealing","method_name":"Cosine Annealing"},{"method_slug":"dense-connections","method_name":"Dense Connections"},{"method_slug":"discriminative-fine-tuning","method_name":"Discriminative Fine-Tuning"},{"method_slug":"dropout","method_name":"Dropout"},{"method_slug":"gpt-2","method_name":"GPT-2"},{"method_slug":"gpt-neo","method_name":"GPT-Neo"},{"method_slug":"layer-normalization","method_name":"Layer Normalization"},{"method_slug":"linear-layer","method_name":"Linear Layer"},{"method_slug":"linear-warmup-with-cosine-annealing","method_name":"Linear Warmup With Cosine Annealing"},{"method_slug":"multi-head-attention","method_name":"Multi-Head Attention"},{"method_slug":"residual-connection","method_name":"Residual Connection"},{"method_slug":"softmax","method_name":"Softmax"},{"method_slug":"test","method_name":"Test"},{"method_slug":"weight-decay","method_name":"Weight Decay"}],"datasets_introduced":[],"methods_introduced":[],"results":[{"leaderboard":"/sota/language-modelling-on-the-pile","task":"Language Modelling","dataset":"The Pile","model":"GPT-2 Large 774M (test-time training on nearest neighbors)","rank_in_archive_order":20,"of":39,"metrics":{"Bits per byte":"0.85"},"uses_additional_data":false}],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2305.18466","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2305.18466"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/socialfoundations/tttlm","reach":{"status":"ok","spdx":"MIT"}}],"summary":{"unverified":5},"by_repo_kind":{"official":{"samples":5,"ran":0,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"00971c933508922f","entry":"build_index","repo":"socialfoundations/tttlm","repo_kind":"official","path":"code/pile_index.py","file_url":"https://github.com/socialfoundations/tttlm/blob/HEAD/code/pile_index.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"00971c933508922f"}},{"code_sha256_prefix":"099b7a975451c130","entry":"build_roberta_index","repo":"socialfoundations/tttlm","repo_kind":"official","path":"code/pile_index.py","file_url":"https://github.com/socialfoundations/tttlm/blob/HEAD/code/pile_index.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"099b7a975451c130"}},{"code_sha256_prefix":"781bd8f051ccee91","entry":"data_to_dict","repo":"socialfoundations/tttlm","repo_kind":"official","path":"code/pile_index.py","file_url":"https://github.com/socialfoundations/tttlm/blob/HEAD/code/pile_index.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"781bd8f051ccee91"}},{"code_sha256_prefix":"2d5424541f361fa4","entry":"get_addresses_from_file","repo":"socialfoundations/tttlm","repo_kind":"official","path":"code/pile_client.py","file_url":"https://github.com/socialfoundations/tttlm/blob/HEAD/code/pile_client.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"2d5424541f361fa4"}},{"code_sha256_prefix":"8bd3239631623911","entry":"probe_servers","repo":"socialfoundations/tttlm","repo_kind":"official","path":"code/pile_client.py","file_url":"https://github.com/socialfoundations/tttlm/blob/HEAD/code/pile_client.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"8bd3239631623911"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}