Super-hard

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking a

Diffusion & Non-Autoregressive Language Models · Problem 4 of 4

Chapter 15Diffusion & Non-Autoregressive Language Models

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking a

Super-hardProblem 4 / 4

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking and compare its output/latency to a greedy AR baseline on a toy vocabulary.

Implement the function/class skeleton in the editor. Any correct approach is accepted.

Hints

solution.pypython

local draft

import time
import torch
import torch.nn.functional as F

def target(L):
    raise NotImplementedError

class Oracle:

    def __call__(self, x):
        raise NotImplementedError

@torch.no_grad()
def diffusion_sample(model, L, steps, mask_id=MASK):
    raise NotImplementedError

@torch.no_grad()
def ar_sample(model, L, mask_id=MASK):
    raise NotImplementedError

⌘/Ctrl + ↵ to submit

AI review

Ready when you are

Submit your solution and a structured review appears here — verdict, score, and concrete feedback. Any correct approach passes.

Chapter 15Diffusion & Non-Autoregressive Language Models

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking a

Super-hardProblem 4 / 4

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking and compare its output/latency to a greedy AR baseline on a toy vocabulary.

Implement the function/class skeleton in the editor. Any correct approach is accepted.

Hints