Module ml::transformer

#include <std/ml/transformer.h> source

Transformer image-generation backbones: DiT (Diffusion Transformer, Peebles & Xie 2023) and JiT ("Just-image- Transformer" — a plain, unconditioned-per-layer transformer backbone, as opposed to DiT's per-layer adaptive conditioning).

Both operate on [seqLen, dModel] token sequences (already-patchified/ embedded — no patchify/positional-encoding logic here, same scope choice attention.h makes) built from std::mha_forward, layer norm (norm.h) and plain feed-forward blocks. Their weights require grad and every op records its backward edges, so they train through trainer.h -- on a GPU too (TrainConfig.device, model_to), where the attention runs as attention.h's fused kernels.

Structs

DiTBlock	Each block computes 6 modulation vectors (shift1,scale1,gate1, shift2,scale2,gate2) from a conditioning vector c via one linear layer, then runs a pre-LN self-attention sublayer and a pre-LN FFN sublayer, each modulated by its half of the 6 and gated additively into the residual stream (the "adaLN-Zero" block from the DiT paper, minus the SiLU the paper puts before WadaLN and minus the zero-init convention — both orthogonal to the block's wiring, which is what this implements and verifies):
JiTBlock	Conditioning is added directly to the token embeddings once, up front (JiT::forward), rather than modulating every block's normalization the way DiT's adaLN does — the architectural distinction this implementation draws between the two. Each block is an ordinary pre-LN transformer sublayer pair:
JiT	A JiT model: the patch embedding and a stack of JiTBlocks. patches: [seqLen, patchDim]; cond: [1, dModel], broadcast-added to every token's embedding once, up front.