sdft-pytorch 0.0.1__tar.gz

sdft_pytorch-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Phil Wang
+
+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.

sdft_pytorch-0.0.1/PKG-INFO

@@ -0,0 +1,39 @@
+Metadata-Version: 2.4
+Name: sdft-pytorch
+Version: 0.0.1
+Summary: SDFT - Pytorch
+Author-email: Phil Wang <lucidrains@gmail.com>
+License: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: discrete-continuous-embed-readout
+Requires-Dist: einops>=0.8.2
+Requires-Dist: ema-pytorch
+Requires-Dist: Jinja2
+Requires-Dist: torch>=2.5
+Requires-Dist: torch-einops-utils>=0.0.21
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: x-transformers; extra == "test"
+Dynamic: license-file
+
+<img src="./sdft.png" width="450px"></img>
+
+## SDFT - Pytorch (wip)
+
+Explorations into the proposed SDFT, [Self-Distillation Enables Continual Learning](https://arxiv.org/abs/2601.19897), from Shenfeld et al. of MIT
+
+## Citations
+
+```bibtex
+@misc{shenfeld2026selfdistillationenablescontinuallearning,
+    title   = {Self-Distillation Enables Continual Learning}, 
+    author  = {Idan Shenfeld and Mehul Damani and Jonas Hübotter and Pulkit Agrawal},
+    year    = {2026},
+    eprint  = {2601.19897},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2601.19897}, 
+}
+```

sdft_pytorch-0.0.1/README.md

@@ -0,0 +1,19 @@
+<img src="./sdft.png" width="450px"></img>
+
+## SDFT - Pytorch (wip)
+
+Explorations into the proposed SDFT, [Self-Distillation Enables Continual Learning](https://arxiv.org/abs/2601.19897), from Shenfeld et al. of MIT
+
+## Citations
+
+```bibtex
+@misc{shenfeld2026selfdistillationenablescontinuallearning,
+    title   = {Self-Distillation Enables Continual Learning}, 
+    author  = {Idan Shenfeld and Mehul Damani and Jonas Hübotter and Pulkit Agrawal},
+    year    = {2026},
+    eprint  = {2601.19897},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2601.19897}, 
+}
+```

sdft_pytorch-0.0.1/pyproject.toml

@@ -0,0 +1,36 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sdft-pytorch"
+version = "0.0.1"
+description = "SDFT - Pytorch"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [
+    { name = "Phil Wang", email = "lucidrains@gmail.com" }
+]
+dependencies = [
+    "discrete-continuous-embed-readout",
+    "einops>=0.8.2",
+    "ema-pytorch",
+    'Jinja2',
+    "torch>=2.5",
+    "torch-einops-utils>=0.0.21"
+]
+
+[project.optional-dependencies]
+test = [
+    "pytest",
+    "x-transformers"
+]
+
+[tool.pytest.ini_options]
+pythonpath = [
+    "."
+]
+
+[tool.setuptools]
+packages = ["sdft_pytorch"]

sdft_pytorch-0.0.1/sdft_pytorch/__init__.py

@@ -0,0 +1,3 @@
+from sdft_pytorch.sdft_pytorch import (
+    SDFT
+)

sdft_pytorch-0.0.1/sdft_pytorch/sdft_pytorch.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+from typing import Callable
+from collections import namedtuple
+
+from jinja2 import Template, Environment, meta
+
+import torch
+from torch.nn import Module
+import torch.nn.functional as F
+from torch import nn, cat, stack, is_tensor, tensor, Tensor
+
+from einops import rearrange
+
+from torch_einops_utils import pad_sequence
+
+from ema_pytorch import EMA
+
+from x_transformers import TransformerWrapper
+
+from discrete_continuous_embed_readout import Readout
+
+# default query / demonstration template for in-context learned distillation targets from teacher for student
+
+DEFAULT_STUDENT_PROMPT_TEMPLATE = """
+[Instruction]
+You are a helpful assistant
+
+[Query]
+{{ question }}
+
+[Response]
+"""
+
+DEFAULT_TEACHER_PROMPT_TEMPLATE = """
+[Task Instructions] You are a helpful assistant. Please answer the question based on the provided logic.
+
+[Expert Demonstration] Question: {{ question }} Expert Reasoning and Answer: {{ answer }}
+
+[Current Task] Question: {{ question }} Answer:
+"""
+
+def get_variables_from_template(template):
+
+    env = Environment()
+
+    parsed_template = env.parse(template)
+
+    return set(meta.find_undeclared_variables(parsed_template))
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def maybe_cast_tensor(t):
+    return t if is_tensor(t) else tensor(t)
+
+# classes
+
+SDFTOutput = namedtuple('SDFTOutput', ('loss', 'response'))
+
+class SDFT(Module):
+    def __init__(
+        self,
+        model: TransformerWrapper,
+        tokenizer_encode: Callable[[list[str]], list[Tensor]],
+        student_max_response_len,
+        student_prompt_template = DEFAULT_STUDENT_PROMPT_TEMPLATE,
+        teacher_update_rate = 0.01,
+        teacher_prompt_template = DEFAULT_TEACHER_PROMPT_TEMPLATE,
+    ):
+        super().__init__()
+
+        self.student = model
+
+        self.teacher = EMA(
+            model,
+            beta = 1. - teacher_update_rate,
+            include_online_model = False
+        )
+
+        # sampling
+
+        self.student_max_response_len = student_max_response_len
+
+        self.discrete_readout = Readout(dim = 0, num_discrete = 1)
+
+        # collection of prompts to list[Int['seq']]
+
+        self.tokenizer_encode = tokenizer_encode
+
+        # store templates
+
+        assert get_variables_from_template(teacher_prompt_template) == {'question', 'answer'}, 'your template must contain only variables `question` and `answer`, embedded like so - {{ question }} ... {{ answer }}'
+        self.teacher_prompt_template = Template(teacher_prompt_template)
+
+        assert get_variables_from_template(student_prompt_template) == {'question'}
+        self.student_prompt_template = Template(student_prompt_template)
+
+    def forward(
+        self,
+        questions: list[str],
+        answers: list[str],
+        student_logit_sample_kwargs: dict = dict()
+    ):
+        encode = self.tokenizer_encode
+        assert len(questions) == len(answers)
+
+        student_vars = [{'question': question} for question in questions]
+        teacher_vars = [{'question': question, 'answer': answer} for question, answer in zip(questions, answers)]
+
+        # ready the prompts for student and teacher
+
+        student_prompts_str = [self.student_prompt_template.render(questions) for questions in student_vars]
+        teacher_prompts_str = [self.teacher_prompt_template.render(question_answers) for question_answers in teacher_vars]
+
+        student_prompt_ids = [maybe_cast_tensor(encode(prompt)) for prompt in student_prompts_str]
+        teacher_prompt_ids = [maybe_cast_tensor(encode(prompt)) for prompt in teacher_prompts_str]
+
+        student_prompt_ids, student_seq_start_pos = pad_sequence(student_prompt_ids, return_lens = True, left = True, pad_lens = True)
+        teacher_prompt_ids, teacher_seq_start_pos = pad_sequence(teacher_prompt_ids, return_lens = True, left = True, pad_lens = True)
+
+        student_cache = None
+        teacher_cache = None
+
+        # accumulate
+
+        student_response = []
+        token_kl_div_losses = []
+
+        for _ in range(self.student_max_response_len):
+
+            # forward for logit of student and teacher
+
+            student_logits, student_cache = self.student(student_prompt_ids, cache = student_cache, seq_start_pos = student_seq_start_pos, return_intermediates = True)
+
+            with torch.no_grad():
+                self.teacher.eval()
+                teacher_logits, teacher_cache = self.teacher(teacher_prompt_ids, cache = teacher_cache, seq_start_pos = teacher_seq_start_pos, return_intermediates = True)
+
+            student_token_logit = student_logits[:, -1:]
+            teacher_token_logit = teacher_logits[:, -1:]
+
+            student_token_log_probs = student_token_logit.log_softmax(dim = -1)
+            teacher_token_probs = teacher_token_logit.softmax(dim = -1)
+
+            # privileged self distillation via ICL
+
+            token_kl_div = F.kl_div(
+                student_token_log_probs,
+                teacher_token_probs,
+                reduction = 'none'
+
+            ).sum(dim = -1)
+
+            token_kl_div_losses.append(token_kl_div)
+
+            # sample
+
+            sampled_action = self.discrete_readout.sample(student_token_logit, **student_logit_sample_kwargs)
+            student_response.append(sampled_action)
+
+            # set student and teacher tokens to the next sampled token
+
+            student_prompt_ids = sampled_action
+            teacher_prompt_ids = sampled_action
+
+        # stack and return
+
+        student_response = cat(student_response, dim = 1)
+        token_kl_div_losses = cat(token_kl_div_losses, dim = 1)
+
+        loss = token_kl_div_losses.mean()
+
+        return SDFTOutput(loss, student_response)

sdft_pytorch-0.0.1/sdft_pytorch.egg-info/PKG-INFO

@@ -0,0 +1,39 @@
+Metadata-Version: 2.4
+Name: sdft-pytorch
+Version: 0.0.1
+Summary: SDFT - Pytorch
+Author-email: Phil Wang <lucidrains@gmail.com>
+License: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: discrete-continuous-embed-readout
+Requires-Dist: einops>=0.8.2
+Requires-Dist: ema-pytorch
+Requires-Dist: Jinja2
+Requires-Dist: torch>=2.5
+Requires-Dist: torch-einops-utils>=0.0.21
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: x-transformers; extra == "test"
+Dynamic: license-file
+
+<img src="./sdft.png" width="450px"></img>
+
+## SDFT - Pytorch (wip)
+
+Explorations into the proposed SDFT, [Self-Distillation Enables Continual Learning](https://arxiv.org/abs/2601.19897), from Shenfeld et al. of MIT
+
+## Citations
+
+```bibtex
+@misc{shenfeld2026selfdistillationenablescontinuallearning,
+    title   = {Self-Distillation Enables Continual Learning}, 
+    author  = {Idan Shenfeld and Mehul Damani and Jonas Hübotter and Pulkit Agrawal},
+    year    = {2026},
+    eprint  = {2601.19897},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2601.19897}, 
+}
+```

sdft_pytorch-0.0.1/sdft_pytorch.egg-info/SOURCES.txt

@@ -0,0 +1,11 @@
+LICENSE
+README.md
+pyproject.toml
+sdft_pytorch/__init__.py
+sdft_pytorch/sdft_pytorch.py
+sdft_pytorch.egg-info/PKG-INFO
+sdft_pytorch.egg-info/SOURCES.txt
+sdft_pytorch.egg-info/dependency_links.txt
+sdft_pytorch.egg-info/requires.txt
+sdft_pytorch.egg-info/top_level.txt
+tests/test_sdft.py

sdft_pytorch-0.0.1/sdft_pytorch.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

sdft_pytorch-0.0.1/sdft_pytorch.egg-info/requires.txt

@@ -0,0 +1,10 @@
+discrete-continuous-embed-readout
+einops>=0.8.2
+ema-pytorch
+Jinja2
+torch>=2.5
+torch-einops-utils>=0.0.21
+
+[test]
+pytest
+x-transformers

sdft_pytorch-0.0.1/sdft_pytorch.egg-info/top_level.txt

@@ -0,0 +1 @@
+sdft_pytorch

sdft_pytorch-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

sdft_pytorch-0.0.1/tests/test_sdft.py

@@ -0,0 +1,36 @@
+import torch
+from sdft_pytorch.sdft_pytorch import SDFT
+
+def test_sdft():
+    from torch import tensor
+    from x_transformers import TransformerWrapper, Decoder
+
+    model = TransformerWrapper(
+        num_tokens = 256,
+        max_seq_len = 512,
+        attn_layers = Decoder(
+            dim = 512,
+            depth = 2
+        )
+    )
+
+    def tokenizer_encode(prompts: list[str]):
+        return [
+            tensor([ord(c) for c in prompt])
+            for prompt in prompts
+        ]
+
+
+    sdft_wrapper = SDFT(
+        model,
+        student_max_response_len = 128,
+        tokenizer_encode = tokenizer_encode,
+    )
+
+    loss, response = sdft_wrapper(
+        questions = ['12+48', '2*3'],
+        answers = ['60', '6']
+    )
+
+    loss.backward()
+    assert response.shape == (2, 128)