canter / TECHNICAL_REPORT.md
data-archetype's picture
Publish Canter 0.5.0 contrastive PDG support
13490a2 verified
|
Raw
History Blame Contribute Delete
22.1 kB

Canter Technical Report

Model card · Example gallery · API and inference parameters

Overview

This is a 2 billion parameter text-to-image model trained from scratch on a single NVIDIA GH200. This was made possible by:

  • Operating in the semantic, diffusion-friendly latent space of DINAC-AE-D2. DINAC-AE-D2 was also trained from scratch for this purpose by the same author and accelerates denoiser convergence.
  • Using SPRINT, which reduces training time and VRAM use while also accelerating convergence.
  • Training on a focused dataset.
  • Using the small SmolLM2-360M language model as the text encoder.
  • Using jagged attention with PyTorch NestedTensor.
  • Using 8-bit AdamW and a CPU-resident exponential moving average to save GPU memory.
  • Patience.

The denoiser generates 128-channel latents with a patch size of 16 by 16 pixels. DINAC-AE-D2 decodes these latents into RGB images. The training data consists primarily of photographic images.

Architecture

Flow matching

This is a flow matching model. As a reminder, let xx be a clean, whitened image latent and let εN(0,I)\varepsilon \sim \mathcal{N}(0, I) be Gaussian noise. Flow matching constructs the linear interpolation path

xt=(1t)x+tε,t[0,1]. x_t = (1-t)x + t\varepsilon, \qquad t \in [0,1].

The clean latent lies at t=0t=0, and pure noise lies at t=1t=1. The target velocity is constant along this path:

v(xt,t)=dxtdt=εx. v^\star(x_t,t) = \frac{\mathrm{d}x_t}{\mathrm{d}t} = \varepsilon - x.

A standard flow-matching model directly predicts this velocity and minimizes

LFM=E[vθ(xt,t)(εx)22]. \mathcal{L}_{\mathrm{FM}} = \mathbb{E}\left[ \left\lVert v_\theta(x_t,t) - (\varepsilon-x) \right\rVert_2^2 \right].

The model retains this flow path, velocity target, and loss while changing the network output parameterization.

xv-pred parameterization

Under the simplifying assumptions that xx and ε\varepsilon are independent and have unit variance, the conditional expectation of the noise given xtx_t is

E[εxt]=tt2+(1t)2xt. \mathbb{E}[\varepsilon \mid x_t] = \frac{t}{t^2 + (1-t)^2}x_t.

Define

f(t)=tt2+(1t)2. f(t) = \frac{t}{t^2 + (1-t)^2}.

Velocity is parameterized as

vθ(xt,t)=uθ(xt,t)+f(t)xt, v_\theta(x_t,t) = u_\theta(x_t,t) + f(t)x_t,

where uθu_\theta is the projected network output. The corresponding network target is

u(xt,t)=(εx)f(t)xt. u^\star(x_t,t) = (\varepsilon-x) - f(t)x_t.

Expanding xtx_t gives

u=(1f(t)t)ε(1+f(t)(1t))x. u^\star = \left(1-f(t)t\right)\varepsilon - \left(1+f(t)(1-t)\right)x.

At t=0t=0, f(0)=0f(0)=0, so u=εxu^\star=\varepsilon-x, the standard velocity target. At t=1t=1, f(1)=1f(1)=1 and xt=εx_t=\varepsilon, so u=xu^\star=-x. This endpoint behavior gives xv-pred its name: it behaves like velocity prediction near the clean endpoint and x-prediction, up to a fixed sign, near the noisy endpoint.

The noise coefficient has the closed form

1f(t)t=(1t)2t2+(1t)2. 1-f(t)t = \frac{(1-t)^2}{t^2+(1-t)^2}.

It approaches zero quadratically as tt approaches 1. In the high-noise region, the analytic f(t)xtf(t)x_t term accounts for the predictable part of the noise. The network can focus on the remaining image-dependent residual. Near the clean endpoint, xv-pred reduces to the original flow-matching target and retains its emphasis on fine detail.

We tested xv-pred against direct velocity prediction and found that it accelerates early training convergence. We retain xv-pred for this reason.

Analytic xv-pred coefficients

The plot shows the magnitudes of the noise and clean-latent coefficients in uu^\star.

Before solver integration, the projected network output and analytic residual are combined in float32 to produce the velocity prediction.

Timestep distribution

Training uses a Beta(2,2)\operatorname{Beta}(2,2) timestep distribution. The usual logit-normal distribution samples the tails too infrequently. We found that this can produce loss spikes during training and poor global image structure at inference.

A resolution-dependent logSNR shift is applied after sampling. During the later high-resolution training stage, the total shift is:

Δ(w,h)=0.9+log(2562wh). \Delta(w,h) = 0.9 + \log\left(\frac{256^2}{wh}\right).

Under this flow convention, a negative shift moves timesteps toward the noisy endpoint at t=1t=1. The shift therefore becomes more negative as image resolution increases. At 102421024^2, the resolution term alone is approximately 2.77-2.77. Once training moved to exclusively high-resolution batches, we added the +0.9+0.9 base shift shown above, giving a total shift of approximately 1.87-1.87. The earlier shift near 3-3 was too aggressive and slowed the convergence of fine image details.

Timesteps are sampled using equal-probability stratification across each optimizer batch. We divide the shifted distribution into one stratum per sample, draw once from every stratum, and randomly permute the assignments. This reduces timestep-sampling variance and is important for stable training at small batch sizes.

High-level design

The denoiser is a 30-layer image transformer with a 3/24/3 layout:

128-channel latent grid
  -> latent and position projection
  -> 3 always-on prefix layers
  -> 24 SPRINT middle layers
  -> timestep-weighted sparse-dense residual fusion
  -> 3 always-on suffix layers
  -> latent velocity projection

All 30 image layers have width 2048, 16 attention heads, and a four-times expansion GELU MLP. The first two prefix layers and last two suffix layers use image self-attention only. The third prefix layer, every middle layer, and the first suffix layer apply text cross-attention before image self-attention.

The prefix and suffix always process the complete image-token grid. During the main SPRINT training path, the 24 middle layers keep one token from each 2 by 2 spatial group and drop the other three. This reduces the middle sequence length by 75 percent.

Spatial encoding

Each 128-channel latent cell is projected to width 2048. Two complementary spatial encodings are then used.

The first is an additive 2D sin/cos embedding. For a latent grid of height HH and width WW, token-center coordinates are normalized by D=max(H,W)D=\max(H,W):

yr=2(r+1/2)D1,xc=2(c+1/2)D1. y_r = \frac{2(r+1/2)}{D}-1, \qquad x_c = \frac{2(c+1/2)}{D}-1.

The embedding contains 512 frequency bands for each of sin(y)\sin(y), cos(y)\cos(y), sin(x)\sin(x), and cos(x)\cos(x). Periods follow 100k/512100^{k/512} for k{0,,511}k \in \{0,\ldots,511\}. The resulting 2048-dimensional features pass through a learned biasless 2048-to-2048 projection in float32 and are added to the latent tokens before the first transformer layer.

The second encoding is axial 2D RoPE in every image self-attention block. Each 128-dimensional attention head is split evenly between the row and column axes, with 32 rotary frequencies per axis, base 10000, and adjacent-channel rotation pairs. RoPE values are constructed in float32 from integer grid coordinates. SPRINT gathers the corresponding row and column rotations when it selects sparse middle tokens, preserving their original positions. The frequency range was critical to image quality in our tests. We retain the older standard base of 10000 because the lower bases used by many recent image models produced substantially worse results.

Image-text cross-attention does not apply 2D RoPE to its image queries. The additive embedding therefore keeps absolute spatial information in the residual stream where it remains available to text cross-attention. Axial RoPE provides the complementary relative geometry used by image self-attention.

Text conditioning

Text conditioning uses the first 24 layers of SmolLM2-360M. Hidden states are taken after layers 8, 16, and 24. Each 960-dimensional tap is projected to width 1024. The three projections are summed in float32 and passed through four trainable text refinement layers.

Each refinement layer uses 8-head self-attention with one-dimensional RoPE, followed by a four-times expansion GELU MLP. The refined tokens provide the keys and values for image-text cross-attention. Image tokens provide the queries.

Cross-attention uses 16 heads with head dimension 128. Timestep-conditioned AdaLN controls the image-query scale and the gated attention residual. The same refined text representation is reused across all denoiser evaluations during sampling.

Text conditioning is dropped for 10 percent of training samples to train the unconditional branch used by classifier-free guidance.

SPRINT routing and residual fusion

Let pp be the dense output of the prefix and let mm be the output of the middle stack. For the sparse path, the retained middle tokens are returned to their original spatial positions and learned mask tokens fill the dropped positions. The model concatenates pp and mm, applies a learned scale derived from the flow-time embedding, and projects the result back to width 2048:

z=Wfuse ⁣([p,m](1+s(SiLU(c(t))))). z = W_{\mathrm{fuse}}\!\left( [p,m] \odot \left(1+s(\operatorname{SiLU}(c(t)))\right) \right).

Here c(t)c(t) is the timestep embedding and ss is a learned linear projection. The time-dependent scale lets the model vary the contribution of the prefix skip and middle path over the flow trajectory. The fused dense token grid then passes through all three suffix layers.

NestedTensor training

Training uses PyTorch jagged NestedTensor attention throughout for two reasons:

  • Text prompts have different lengths. Packing their tokens avoids padding and its associated attention cost. It also avoids exposure bias from padded training, where the model can learn to use padding positions as register-like scratch space.
  • Image samples can use different SPRINT paths within the same mini-batch. The jagged middle stack processes dense and sparse token sequences together without padding either sequence.

The SPRINT path mix is:

Fraction Middle path
10% Complete middle-path drop
10% All middle layers with no token drop
80% All middle layers with 75% token drop

The 10 percent dense path is needed to limit exposure bias from training the middle stack only on sparse token grids. The original SPRINT paper also found that fully sparse training reduced quality and addressed this with a short dense post-training stage. We found that dense post-training was detrimental. Mixing 10 percent dense-path samples throughout training gave the best results.

We interpret this result in two ways:

  • Sparse middle layers place the model in a prefix-conditioned, mask-modelling regime. This increases image coherence.
  • The train-to-inference discrepancy weakens the full main path. Without this effect, the main path becomes too strong relative to the dropped path used for path-drop guidance, producing excessive sharpening and contrast.

Architecture summary

Component Setting
Latent representation DINAC-AE-D2, 128 channels, spatial stride 16
Image width 2048
Image layers 30 total: 3 prefix, 24 middle, 3 suffix
Image attention 16 heads, head dimension 128
Image position encoding Learned projection of normalized 2D sin/cos features and axial 2D RoPE
Image MLP GELU, expansion ratio 4
Image conditioning Shared AdaLN base with rank-256 per-layer deltas
Text backbone First 24 SmolLM2-360M layers
Text taps Layers 8, 16, and 24
Refined text width 1024
Text refinement 4 layers, 8 attention heads, 1D RoPE
Image-text cross-attention 16 heads, head dimension 128
Cross-attention placement Prefix layer 3, all 24 middle layers, suffix layer 1

Image DiT block

Every image layer contains a self-attention residual followed by an MLP residual. Timestep conditioning comes from a shared AdaLN projection plus a rank-256 delta for each block:

image tokens [B, N, 2048]
  -> RMSNorm
  -> timestep AdaLN scale
  -> biasless QKV projection
  -> per-head RMSNorm on Q and K
  -> axial 2D RoPE on Q and K
  -> scaled dot-product attention
  -> biasless output projection
  -> RMSNorm
  -> tanh timestep gate
  -> residual add
  -> RMSNorm
  -> timestep AdaLN scale
  -> biasless Linear(2048 -> 8192)
  -> GELU
  -> biasless Linear(8192 -> 2048)
  -> RMSNorm
  -> tanh timestep gate
  -> residual add

Text refinement block

The four text refinement layers use a pre-norm transformer block:

text tokens [total_tokens, 1024]
  -> RMSNorm
  -> biasless QKV projection
  -> per-head RMSNorm on Q and K
  -> 1D RoPE on Q and K
  -> 8-head scaled dot-product attention
  -> biasless output projection
  -> residual add
  -> RMSNorm
  -> biasless Linear(1024 -> 4096)
  -> GELU
  -> biasless Linear(4096 -> 1024)
  -> residual add

The jagged layout stores only valid prompt tokens. The same block can also use a dense masked attention backend for inference.

Image-text cross-attention block

Cross-conditioned image layers apply this residual before their image DiT block:

image tokens [B, N, 2048]
  -> RMSNorm
  -> timestep AdaLN scale
  -> biasless query projection

text tokens [total_tokens, 1024]
  -> RMSNorm
  -> biasless key-value projection

queries, keys, values
  -> 16 heads of width 128
  -> per-head RMSNorm on queries and keys
  -> scaled dot-product attention
  -> biasless Linear(2048 -> 2048)
  -> tanh timestep gate
  -> image residual add

The third prefix layer and first suffix layer use learned scalar multipliers on their cross-attention and image-block timestep modulation.

Inference

Path-drop guidance (PDG) is used by default. Following SPRINT, the weak prediction skips all 24 middle layers while the main prediction uses the full model:

vPDG=vweak+sPDG(vmainvweak). v_{\mathrm{PDG}} = v_{\mathrm{weak}} + s_{\mathrm{PDG}}\left(v_{\mathrm{main}}-v_{\mathrm{weak}}\right).

PDG produces substantially better images than classifier-free guidance (CFG) for this model. It is also cheaper because its weak path evaluates only the six always-on image layers.

The guidance is strong. High PDG scales can cause excessive sharpening, contrast, and structural defects. We mitigate this with self_attention_gain, applied only to image self-attention on the main path. For gain gg, every image self-attention query is scaled by exp(g)\exp(g):

Attentiong(Q,K,V)=softmax(exp(g)QKTd)V. \operatorname{Attention}_g(Q,K,V) = \operatorname{softmax}\left( \frac{\exp(g)QK^\mathsf{T}}{\sqrt{d}} \right)V.

This is equivalent to a softmax temperature T=exp(g)T=\exp(-g). A negative gain therefore raises the effective temperature, softens main-path attention, and empirically softens the main-path image distribution.

At resolutions around the 102421024^2 aspect-ratio buckets, a PDG scale near 2.5 and a self-attention gain near -0.03 gives substantially better results than reducing PDG while leaving the gain at zero.

Contrastive PDG conditions the middle-skipped path with either a negative prompt or Canter's learned unconditional text. PDG can also be combined with CFG. The release provides several CFG/PDG interaction modes, constant, linear, and power PDG curves, and independent start and stop steps for both guidance methods. See API and inference parameters for the complete interface.

Dataset

Image collection and captions

Training uses about 17 million publicly available images. Around 4.8 million of these form a higher-quality, high-resolution subset used for late-stage training. The data pipeline includes repeated image deduplication together with image and caption quality checks.

Images were captioned with Gemini, ChatGPT, Qwen, and Mistral models, with Gemini providing most of the captions.

Aspect-ratio buckets and resolution curriculum

Training uses SDXL-style aspect-ratio buckets. Each nominal resolution defines a fixed set of landscape, square, and portrait shapes scaled from the SDXL 1024-pixel bucket table. Width and height are divisible by 32. Each image is assigned to the bucket that removes the least content, resized to cover that shape while preserving its aspect ratio, and cropped to the final dimensions. Every mini-batch contains images from one bucket, so no spatial padding is needed.

Resolution was progressively increased from nominal 2562256^2 aspect-ratio buckets to 102421024^2 buckets. Intermediate stages introduced 3842384^2, 5122512^2, and 7682768^2 bucket families before the final 102421024^2 stage. The nominal resolution describes the scale of the bucket family rather than requiring square images.

Synthetic text data

One percent of training samples use procedurally rendered text images paired with captions that state the text shown in the image. This augmentation was introduced in two stages:

  1. The first stage renders individual letters, digits, symbols, and common words in high-contrast sans-serif and serif fonts. It varies case, size, style, color, and position.
  2. The second stage broadens the text distribution with a filtered English vocabulary, short phrases, and punctuation- or digit-heavy strings designed to stress tokenization. It also adds occasional rotated text.

The generated captions always preserve the rendered string exactly while varying descriptions of its font, case, style, color, size, and location.

Caption augmentation

Training applies a small set of lightweight text augmentations to the source captions. These include lowercasing ordinary words, dropping terminal punctuation, replacing one space with a newline or double space, and changing spacing around a comma. Quoted spans are protected so text that should appear verbatim in an image is not modified.

Training

Training uses custom PyTorch training code. Model weights remain in float32. Forward and backward passes run under bfloat16 CUDA autocast, with explicit float32 computation retained for numerically sensitive operations.

Block compilation

The transformer is compiled block by block with torch.compile rather than as one graph. The dense, always-on image blocks are captured as full graphs. SPRINT middle blocks, text refinement blocks, and image-text cross-attention use dynamic, non-fullgraph compilation to support variable image and text sequence lengths with jagged NestedTensor. SPRINT routing and block sequencing remain outside the compiled blocks. Training and evaluation use separate compiled callables.

Optimizer and EMA

Optimization uses torchao.optim.AdamW8bit with a compiled optimizer step. We tested this optimizer against standard AdamW and observed no difference in convergence. The first-moment coefficient is β1=0.9\beta_1=0.9, epsilon is 10810^{-8}, and weight decay is zero.

An exponential moving average of the model weights is stored in float32 on the CPU and updated every 10 optimizer steps. The GPU-to-CPU snapshot and CPU update are asynchronous. The decay applied at each update is 0.9999100.9999^{10}, preserving the time scale of a per-step decay of 0.99990.9999 while removing the persistent EMA copy from GPU memory.

Batch size and optimizer scaling

The reference optimizer settings use a batch size of 128, a learning rate of 10410^{-4}, and β2=0.98\beta_2=0.98. For optimizer batch size BB, the learning rate is scaled linearly:

η(B)=104B128. \eta(B) = 10^{-4}\frac{B}{128}.

The second-moment coefficient follows a power rule:

β2(B)=0.98B/128. \beta_2(B) = 0.98^{B/128}.

This keeps the second-moment time scale approximately constant when measured in training samples rather than optimizer steps.

Training began with batch size 128 at the nominal 2562256^2 aspect-ratio buckets. The batch size was progressively reduced as resolution increased, reaching 12 for the final stage using only 102421024^2 aspect-ratio buckets. At batch size 12, the formulas give a learning rate of 9.375×1069.375\times10^{-6} and β20.9981\beta_2\approx0.9981; the stage uses the rounded values 10510^{-5} and 0.9980.998. The learning rate has a 4,000-step linear warmup and remains constant afterward.

Citation

@misc{canter,
  title   = {Canter: An Efficient, Photography-Oriented Text-to-Image Model},
  author  = {data-archetype},
  email   = {data-archetype@proton.me},
  year    = {2026},
  month   = jul,
  url     = {https://huggingface.co/data-archetype/canter},
}

References

DINAC-AE-D2

data-archetype, “DINAC-AE-D2: a DINOv2-aligned class-token diffusion autoencoder,” 2026. Technical report.

Flow Matching

Lipman et al., “Flow Matching for Generative Modeling,” ICLR 2023. arXiv:2210.02747.

SDXL

Podell et al., “SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis,” ICLR 2024. arXiv:2307.01952.

SPRINT

Park et al., “Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers,” 2025. arXiv:2510.21986.

SmolLM2

Allal et al., “SmolLM2: When Smol Goes Big: Data-Centric Training of a Small Language Model,” 2025. arXiv:2502.02737. SmolLM2-360M model card.

TorchAO AdamW8bit

PyTorch, “PyTorch Native Architecture Optimization: torchao,” 2024. TorchAO overview. AdamW8bit implementation.

Rotary Position Embedding

Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding,” 2021. arXiv:2104.09864.

Classifier-free guidance

Ho and Salimans, “Classifier-Free Diffusion Guidance,” 2022. arXiv:2207.12598.

AdamW

Loshchilov and Hutter, “Decoupled Weight Decay Regularization,” ICLR 2019. arXiv:1711.05101.