TECHNICAL_REPORT.md

# Canter Technical Report

[Model card](README.md) · [Example gallery](GALLERY.md) ·
[API and inference parameters](API.md)

## 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](https://huggingface.co/data-archetype/dinac_ae_d2).
  DINAC-AE-D2 was also trained from scratch for this purpose by the same author and accelerates
  denoiser convergence.
- Using [SPRINT](#sprint), which reduces training time and VRAM use while also
  accelerating convergence.
- Training on a focused dataset.
- Using the small
  [SmolLM2-360M](https://huggingface.co/HuggingFaceTB/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 \\(x\\) be a clean, whitened
image latent and let \\(\varepsilon \sim \mathcal{N}(0, I)\\) be Gaussian noise.
Flow matching constructs the linear interpolation path

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

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

$$
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

$$
\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 \\(x\\) and \\(\varepsilon\\) are
independent and have unit variance, the conditional expectation of the noise
given \\(x_t\\) is

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

Define

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

Velocity is parameterized as

$$
v_\theta(x_t,t) = u_\theta(x_t,t) + f(t)x_t,
$$

where \\(u_\theta\\) is the projected network output. The corresponding network
target is

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

Expanding \\(x_t\\) gives

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

At \\(t=0\\), \\(f(0)=0\\), so \\(u^\star=\varepsilon-x\\), the standard velocity
target. At \\(t=1\\), \\(f(1)=1\\) and \\(x_t=\varepsilon\\), so \\(u^\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

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

It approaches zero quadratically as \\(t\\) approaches 1. In the high-noise
region, the analytic \\(f(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](assets/xv_pred_coefficients.png)

The plot shows the magnitudes of the noise and clean-latent coefficients in
\\(u^\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 \\(\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:

$$
\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=1\\). The shift therefore becomes more negative as image
resolution increases. At \\(1024^2\\), the resolution term alone is approximately
\\(-2.77\\). Once training moved to exclusively high-resolution batches, we added
the \\(+0.9\\) base shift shown above, giving a total shift of approximately
\\(-1.87\\). The earlier shift near \\(-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:

```text
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 \\(H\\)
and width \\(W\\), token-center coordinates are normalized by
\\(D=\max(H,W)\\):

$$
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)\\), \\(\cos(y)\\),
\\(\sin(x)\\), and \\(\cos(x)\\). Periods follow
\\(100^{k/512}\\) for \\(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 \\(p\\) be the dense output of the prefix and let \\(m\\) 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 \\(p\\) and \\(m\\), applies a learned scale derived
from the flow-time embedding, and projects the result back to width 2048:

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

Here \\(c(t)\\) is the timestep embedding and \\(s\\) 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:

```text
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
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:

```text
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:

$$
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 \\(g\\), every image self-attention query is scaled by \\(\exp(g)\\):

$$
\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)\\). A negative gain
therefore raises the effective temperature, softens main-path attention, and
empirically softens the main-path image distribution.

At resolutions around the \\(1024^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](API.md) 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 \\(256^2\\) aspect-ratio
buckets to \\(1024^2\\) buckets. Intermediate stages introduced \\(384^2\\), \\(512^2\\),
and \\(768^2\\) bucket families before the final \\(1024^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`](https://github.com/pytorch/ao/blob/main/torchao/optim/adam.py)
with a compiled optimizer step. We tested this optimizer against standard
AdamW and observed no difference in convergence. The first-moment coefficient
is \\(\beta_1=0.9\\), epsilon is \\(10^{-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.9999^{10}\\), preserving the time scale of a per-step decay of \\(0.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
\\(10^{-4}\\), and \\(\beta_2=0.98\\). For optimizer batch size \\(B\\), the learning rate
is scaled linearly:

$$
\eta(B) = 10^{-4}\frac{B}{128}.
$$

The second-moment coefficient follows a power rule:

$$
\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 \\(256^2\\) aspect-ratio
buckets. The batch size was progressively reduced as resolution increased,
reaching 12 for the final stage using only \\(1024^2\\) aspect-ratio buckets. At
batch size 12, the formulas give a learning rate of \\(9.375\times10^{-6}\\) and
\\(\beta_2\approx0.9981\\); the stage uses the rounded values \\(10^{-5}\\) and
\\(0.998\\). The learning rate has a 4,000-step linear warmup and remains constant
afterward.

## Citation

```bibtex
@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](https://huggingface.co/data-archetype/dinac_ae_d2).

### Flow Matching

Lipman et al., “Flow Matching for Generative Modeling,” ICLR 2023.
[arXiv:2210.02747](https://arxiv.org/abs/2210.02747).

### SDXL

Podell et al., “SDXL: Improving Latent Diffusion Models for High-Resolution
Image Synthesis,” ICLR 2024.
[arXiv:2307.01952](https://arxiv.org/abs/2307.01952).

### SPRINT

Park et al., “Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion
Transformers,” 2025. [arXiv:2510.21986](https://arxiv.org/abs/2510.21986).

### SmolLM2

Allal et al., “SmolLM2: When Smol Goes Big: Data-Centric Training of a Small
Language Model,” 2025.
[arXiv:2502.02737](https://arxiv.org/abs/2502.02737).
[SmolLM2-360M model card](https://huggingface.co/HuggingFaceTB/SmolLM2-360M).

### TorchAO AdamW8bit

PyTorch, “PyTorch Native Architecture Optimization: torchao,” 2024.
[TorchAO overview](https://pytorch.org/blog/pytorch-native-architecture-optimization/).
[AdamW8bit implementation](https://github.com/pytorch/ao/blob/main/torchao/optim/adam.py).

### Rotary Position Embedding

Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding,”
2021. [arXiv:2104.09864](https://arxiv.org/abs/2104.09864).

### Classifier-free guidance

Ho and Salimans, “Classifier-Free Diffusion Guidance,” 2022.
[arXiv:2207.12598](https://arxiv.org/abs/2207.12598).

### AdamW

Loshchilov and Hutter, “Decoupled Weight Decay Regularization,” ICLR 2019.
[arXiv:1711.05101](https://arxiv.org/abs/1711.05101).