PTQ-AttnDM, Explained Simply

This post is a shorter and more readable version of my dissertation, PTQ-AttnDM: Self-Attention Enhanced Post-Training Quantisation for Diffusion Models.

The core question behind the work is straightforward:

Can we make diffusion models cheaper to run without retraining them from scratch and without destroying image quality?

My answer was partly yes. The strongest result came from treating self-attention more carefully during quantisation instead of quantising every part of the model in the same way.

Why this problem matters

Diffusion models are powerful, but they are expensive to deploy:

• they run through many denoising steps
• they use large neural networks with heavy memory and compute cost
• they are difficult to use on smaller or cheaper hardware

Quantisation is an attractive solution because it reduces the precision of weights and activations, which can reduce memory use and speed up inference. The challenge is that diffusion models are unusually sensitive to quantisation noise.

That sensitivity becomes more obvious in self-attention, where small numerical errors can compound across timesteps and damage generation quality.

The key idea

Instead of applying one uniform quantisation rule everywhere, PTQ-AttnDM focuses on where diffusion models are most fragile.

The method concentrates on self-attention and asks three practical questions:

1. Which parts of attention are most sensitive to precision loss?
2. How does that sensitivity change over diffusion timesteps?
3. Can we spend precision only where it matters most?

The system is designed to answer those questions during calibration, then apply quantisation more selectively during inference.

Easy visual representation

What PTQ-AttnDM changes

1. Mixed precision inside self-attention

The model does not treat query, key, value, and output projections as equally important.

In the implementation, the default pattern is:

• query: higher precision
• key: slightly lower precision
• value: higher precision
• output: higher precision

The reasoning is simple. Query and value contribute more directly to the final representation, while key can often tolerate a little more compression.

2. Group-wise quantisation

Attention channels are grouped instead of being quantised as one large block.

This matters because attention heads can have different activation ranges. If everything shares one quantisation range, stronger channels can distort weaker ones. Grouping makes the quantisation more local and more stable.

3. Multi-timestep calibration

A major difficulty in diffusion PTQ is that activation distributions are not stable across timesteps.

So the method calibrates using multiple denoising steps instead of pretending one calibration snapshot is enough. This is important because the same layer can behave very differently at different points in the generation process.

4. Attention-aware precision scheduling

PTQ-AttnDM also explores adjusting effective precision based on timestep importance.

The goal is to spend more precision when attention is more sensitive and relax it when compression is safer. Conceptually, it is a way of saying:

not every denoising step deserves the same numerical budget

What I evaluated

The paper compares PTQ-AttnDM against other PTQ baselines on diffusion generation quality, mainly using:

• FID
• sFID
• Inception Score

It also includes qualitative image comparisons and ablation studies to test which parts of the method were actually helpful.

Paper artifacts and highlights

Here are a few concrete fragments from the dissertation that show what the method actually looked like in practice.

PTQ-AttnDM algorithm sketch

The high-level algorithm in the paper looks like this:

1. Identify attention and convolution layers
2. Measure attention sensitivity for q, k, v, and output
3. Allocate bits per projection
4. Optimise grouping per projection
5. Quantise attention layers with the chosen config
6. Quantise convolution layers with standard PTQ
7. Assemble the quantised model
8. Calibrate on representative diffusion timesteps
9. Run quantised attention forward with group-wise quantisation

Bit allocation inside attention

One of the most important implementation choices was to avoid giving every attention projection the same precision:

self.bit_config = {
  "query": args.bitwidth if args else 8,
  "key": max(4, args.bitwidth - 2) if args else 6,
  "value": args.bitwidth if args else 8,
  "output": args.bitwidth if args else 8,
}

Quantised attention block construction

The paper also implements a custom attention constructor that turns mixed precision on by default:

def create_enhanced_attention(in_channels, sequence, args):
  return EnhancedQSelfAttention(
    in_channels,
    quantization=True,
    sequence=sequence,
    args=args,
    mixed_precision=True,
    bit_config={
      "query": args.bitwidth,
      "key": max(4, args.bitwidth - 2),
      "value": args.bitwidth,
      "output": args.bitwidth,
    },
  )

Calibration hook snippet

The calibration stage in the paper explicitly captures the min and max range of attention scores:

def qk_hook(module, input, output):
  qk_mins.append(output.min().item())
  qk_maxs.append(output.max().item())

for module in self.attention_modules:
  hooks.append(module.register_forward_hook(qk_hook))

_ = self.model(sample_batch, t=t_tensor)

for hook in hooks:
  hook.remove()

Timestep-aware precision scheduling

Another important artifact is the scheduling rule used to vary effective precision over time:

effective_bits = base_bit_width + 2 * sigmoid(timestep_importance[t])

Sensitivity measurement

To track which layers were fragile, the paper defined a simple sensitivity score from activation statistics:

Sensitivity_L = Range_L * (1 + Variance_L / Mean_L)

Citation highlights

The paper is positioned against a few important reference points:

Results snapshot from the paper

The most important result snippet to carry into the blog is the one that shows both the win and the limitation:

250 steps, W8A8
Q-Diffusion   FID 8.07
APQ-DM        FID 5.51
PTQ-AttnDM    FID 4.43

250 steps, W4A4
Q-Diffusion   FID 13.81
APQ-DM        FID 21.21
PTQ-AttnDM    FID 28.99

Ablation snippet

The grouping ablation is one of the clearest internal signals in the paper:

W8A8, FID by number of groups
1 group   5.79
4 groups  5.34
8 groups  4.95
16 groups 5.05

Main results

The results are mixed, which is exactly why they are useful.

Where the method worked well

In one of the strongest settings, 8-bit weights and activations with 250 timesteps, PTQ-AttnDM achieved:

• FID 4.43
• compared with APQ-DM at 5.51

That means the method was competitive and in that setting actually outperformed a strong baseline.

The ablation results also support two design choices:

• 8-bit attention is much safer than 4-bit attention
• group-wise quantisation helps, with around 8 groups giving the best trade-off in the reported experiments

The calibration study also showed that active timestep sampling was better than random or heuristic sampling, reaching FID 5.49 at W8A8 with a larger calibration set.

Where the method struggled

The method was not consistently best across all settings.

The biggest weakness appeared in more aggressive low-bit settings such as W4A4, where image quality dropped sharply. In other words, the method helped, but it did not solve the hardest part of diffusion PTQ.

That limitation is important. It suggests that:

• attention-aware PTQ helps most at moderate compression
• very low-bit diffusion inference still needs better techniques
• calibration and scheduling are still fragile when precision becomes too low

What I think the paper really shows

The main contribution is not just a single score. It is a more useful design lesson:

self-attention should not be treated as just another layer during diffusion model quantisation

That insight shaped the whole project.

The paper shows that diffusion quantisation works better when we:

• isolate sensitive attention components
• calibrate across timesteps instead of using one static view
• apply precision selectively rather than uniformly

Even where PTQ-AttnDM does not win outright, it helps explain why diffusion PTQ is hard and where future improvements should focus.

Short version of the methodology

If I had to explain the method in plain English to someone outside ML, I would put it like this:

1. Take a trained diffusion model.
2. Find the parts that are most likely to break when numbers are compressed.
3. Compress those parts more carefully than everything else.
4. Check the model across different denoising stages, because the same layer behaves differently over time.
5. Run the smaller, cheaper model and compare the image quality with full precision and other PTQ methods.

That is the whole idea.

Limitations

A fair reading of the paper should include the limits:

• the gains are strongest in selected 8-bit settings, not everywhere
• 4-bit quantisation remains unstable
• diffusion models are highly timestep-dependent, which makes calibration hard
• compute limits constrained how much experimentation could be done

So this is best understood as a targeted improvement, not a complete solution to diffusion compression.

Future work

The next steps that seem most promising are:

• comparing PTQ behaviour more directly between UNet-based diffusion models and diffusion transformers
• improving low-bit attention quantisation beyond 8-bit
• learning better timestep importance schedules
• finding stronger calibration sets for diffusion-specific activation behaviour

As larger generative models become more common, these compression questions will matter even more.

Final takeaway

PTQ-AttnDM argues for a simple position:

If you want post-training quantisation to work for diffusion models, pay special attention to self-attention.

That does not magically make low-bit diffusion easy, but it does move the problem in a more useful direction.

If you want the full code behind the project, it is available here:

PTQ-AttnDM on GitHub