"""Original reduced DiT backbone with adaLN-Zero and two output branches. A 2-channel 4x4 latent -> four 2x2 patches -> width-eight Transformer -> 2-channel epsilon plus 2 raw variance-interpolation channels. No VAE or sampler. """ import torch import torch.nn as nn import torch.nn.functional as F def positions_2d(): y, x = torch.meshgrid(torch.arange(2), torch.arange(2), indexing="ij") freq = torch.tensor([1.0, 0.01]) px = x.flatten()[:, None] * freq py = y.flatten()[:, None] * freq return torch.cat([px.sin(), px.cos(), py.sin(), py.cos()], -1)[None] class TimestepEmbedding(nn.Module): def __init__(self): super().__init__() self.up = nn.Linear(8, 8) self.silu = nn.SiLU() self.down = nn.Linear(8, 8) def forward(self, t): freq = torch.exp( -torch.arange(4, device=t.device, dtype=self.up.weight.dtype) * torch.log(torch.tensor(10000.0, device=t.device)) / 4 ) angles = t[:, None] * freq return self.down( self.silu(self.up(torch.cat([angles.cos(), angles.sin()], -1))) ) class SelfAttention(nn.Module): def __init__(self): super().__init__() self.query = nn.Linear(8, 8) self.key = nn.Linear(8, 8) self.value = nn.Linear(8, 8) self.softmax = nn.Softmax(-1) self.project = nn.Linear(8, 8) def forward(self, x): b, t, _ = x.shape q = self.query(x).reshape(b, t, 2, 4).transpose(1, 2) k = self.key(x).reshape(b, t, 2, 4).transpose(1, 2) v = self.value(x).reshape(b, t, 2, 4).transpose(1, 2) weights = self.softmax(q @ k.transpose(-1, -2) / 2) return self.project((weights @ v).transpose(1, 2).reshape(b, t, 8)) class AdaLNZeroBlock(nn.Module): def __init__(self): super().__init__() self.silu = nn.SiLU() self.modulation = nn.Linear(8, 48) self.attention_norm = nn.LayerNorm(8, elementwise_affine=False) self.attention = SelfAttention() self.ff_norm = nn.LayerNorm(8, elementwise_affine=False) self.up = nn.Linear(8, 32) self.gelu = nn.GELU(approximate="tanh") self.down = nn.Linear(32, 8) nn.init.zeros_(self.modulation.weight) nn.init.zeros_(self.modulation.bias) def forward(self, x, condition): shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = self.modulation( self.silu(condition) ).chunk(6, -1) h = self.attention_norm(x) * (1 + scale_a[:, None]) + shift_a[:, None] x = x + gate_a[:, None] * self.attention(h) h = self.ff_norm(x) * (1 + scale_f[:, None]) + shift_f[:, None] return x + gate_f[:, None] * self.down(self.gelu(self.up(h))) class PatchDecoder(nn.Module): def __init__(self): super().__init__() self.norm = nn.LayerNorm(8, elementwise_affine=False) self.modulation = nn.Linear(8, 16) self.project = nn.Linear(8, 16) nn.init.zeros_(self.modulation.weight) nn.init.zeros_(self.modulation.bias) nn.init.zeros_(self.project.weight) nn.init.zeros_(self.project.bias) def forward(self, x, condition): shift, scale = self.modulation(F.silu(condition)).chunk(2, -1) patches = self.project(self.norm(x) * (1 + scale[:, None]) + shift[:, None]) return unpatchify(patches) def unpatchify(patches): # Row-major patch positions, then row/column within each patch, then channel. b = patches.shape[0] return ( patches.reshape(b, 2, 2, 2, 2, 4).permute(0, 5, 1, 3, 2, 4).reshape(b, 4, 4, 4) ) class TinyDiT(nn.Module): def __init__(self): super().__init__() self.patchify = nn.Conv2d(2, 8, 2, stride=2) self.register_buffer("positions", positions_2d()) self.time_embedding = TimestepEmbedding() self.class_embedding = nn.Embedding(4, 8) self.block = AdaLNZeroBlock() self.decoder = PatchDecoder() def forward(self, latent, timestep, label): x = self.patchify(latent).flatten(2).transpose(1, 2) + self.positions condition = self.time_embedding(timestep) + self.class_embedding(label) return self.decoder(self.block(x, condition), condition) def setup(): torch.manual_seed(47) model = TinyDiT().double() latent = torch.randn(1, 2, 4, 4, dtype=torch.float64) return model, latent def activate_for_comparison(model): # Deliberate nonzero parameter intervention, never a trained checkpoint. torch.manual_seed(53) with torch.no_grad(): model.block.modulation.weight.normal_(0, 0.08) model.decoder.modulation.weight.normal_(0, 0.08) model.decoder.project.weight.normal_(0, 0.1) def verify(): model, latent = setup() t = torch.tensor([6]) label = torch.tensor([1]) patches = model.patchify(latent).flatten(2).transpose(1, 2) reference = ( F.unfold(latent, 2, stride=2).transpose(1, 2) @ model.patchify.weight.flatten(1).T + model.patchify.bias ) torch.testing.assert_close(patches, reference) x = patches + model.positions c = model.time_embedding(t) + model.class_embedding(label) torch.testing.assert_close(model.block(x, c), x, atol=0, rtol=0) assert torch.count_nonzero(model(latent, t, label)) == 0 ids = torch.arange(64.0).reshape(1, 4, 16) spatial = unpatchify(ids) for patch in range(4): row, col = divmod(patch, 2) torch.testing.assert_close( spatial[0, :, row * 2 : row * 2 + 2, col * 2 : col * 2 + 2] .permute(1, 2, 0) .flatten(), ids[0, patch], ) loss = F.mse_loss(model(latent, t, label)[:, :2], torch.ones_like(latent)) loss.backward() assert model.decoder.project.weight.grad.norm() > 0 assert torch.count_nonzero(model.block.modulation.weight.grad) == 0 activate_for_comparison(model) first = model(latent, t, label) assert not torch.allclose(first, model(latent, t, torch.tensor([2]))) assert not torch.allclose(first, model(latent, torch.tensor([11]), label)) return [ "Convolutional patch embeddings equal explicit unfold plus linear projection", "Zero modulation initializes the residual block as identity and the final prediction as zero", "Unpatchification preserves every patch coordinate and output channel", "Initial gradients reach the output projection; after a declared nonzero intervention both timestep and class change predictions", ] def experiment(): model, latent = setup() cases = [] with torch.no_grad(): for identity, label, t, cls in [ ("zero", "Zero initialization", 6, 1), ("open", "Nonzero parameter intervention", 6, 1), ("class", "Same intervention · class 2", 6, 2), ("time", "Same intervention · timestep 11", 11, 1), ]: if identity == "open": activate_for_comparison(model) time = torch.tensor([t]) category = torch.tensor([cls]) x = model.patchify(latent).flatten(2).transpose(1, 2) + model.positions c = model.time_embedding(time) + model.class_embedding(category) hidden = model.block(x, c) out = model(latent, time, category) mod = model.block.modulation(model.block.silu(c)).reshape(6, 8) cases.append( dict( id=identity, label=label, target="block.modulation", note=( "At initialization all modulation and final projection weights are zero. The block is exactly identity and both output branches are zero. " if identity == "zero" else "A seeded nonzero parameter intervention opens the block and output projection. It is not training or a released DiT checkpoint. " ) + "The first two spatial channels predict epsilon; the last two are raw variance-interpolation parameters, not positive variances. Class and timestep controls hold the latent fixed.", matrices=[ dict( label="Input latent · channel 0", values=latent[0, 0].tolist(), ), dict( label="Attention / FFN shift, scale and gate · six rows", values=mod.tolist(), ), dict( label="Predicted epsilon · channel 0", values=out[0, 0].tolist(), ), dict( label="Raw variance parameter · channel 0", values=out[0, 2].tolist(), ), ], vectors=[dict(label="Condition embedding", values=c[0].tolist())], metrics=[ dict(label="Timestep", value=t), dict(label="Class ID", value=cls), dict( label="Maximum block change from input", value=(hidden - x).abs().max().item(), ), dict(label="Latent patch tokens", value=4), ], ) ) return dict( kind="matrices", title="Open the residual gates with conditioning.", description="Recorded adaLN-Zero initialization and a controlled nonzero parameter intervention. A Transformer maps latent patches back to spatial noise and variance parameters.", controlLabel="DiT condition", cases=cases, )