Direct Alignment: DPO, Rejection Sampling & Distillation4 / 5
Direct Alignment: DPO, Rejection Sampling & Distillation
Implement on-policy distillation: sample from the student, score tokens under a (toy) teacher, and apply the reverse-KL policy-gradient update.
class TinyLMComplete the starter code in the editor. Any correct approach is accepted.
import torch
import torch.nn as nn
import torch.nn.functional as F
class TinyLM(nn.Module):
def __init__(self, V, ctx=4):
raise NotImplementedError
def step_logits(self, last_tok):
raise NotImplementedError
@torch.no_grad()
def rollout(student, V, B, T, start=0):
raise NotImplementedError
def distill_step(student, teacher, opt, V, B=64, T=6, gamma=0.99):
raise NotImplementedError⌘/Ctrl↵to submit
Submit your solution for a score and feedback on correctness, edge cases, and clarity. Any correct approach is accepted.
Direct Alignment: DPO, Rejection Sampling & Distillation
Implement on-policy distillation: sample from the student, score tokens under a (toy) teacher, and apply the reverse-KL policy-gradient update.
class TinyLMComplete the starter code in the editor. Any correct approach is accepted.