BADiff: Bandwidth Adaptive Diffusion Model

[W1, disconnection between different components]

We agree that our original presentation could be improved to clarify the interactions between different model components. We will revise Sections 3.3–3.5 to make these connections more explicit, explain the overall training objective more clearly, and cleanly separate the training and inference phases. Specifically:

Clarifying relationships between components:
The entropy hinge loss (Section 3.3) encourages the model to produce images whose estimated entropy falls within a desired target range. The policy network (Section 3.4) learns to predict the earliest timestep at which this entropy constraint is satisfied. These two components are thus aligned in their objective: to guide the sampling process toward entropy-constrained generation. We will explicitly describe this relationship in the revision.

Unified interpretation of the loss function (Equation 16):
The four terms in Equation 16 are designed to jointly optimize different aspects of the system:

• the diffusion loss ensures visual fidelity,
• the entropy hinge loss enforces bitrate constraints,
• the calibration loss encourages consistency with external codecs,
• and the policy loss trains the adaptive stopping mechanism.

Together, they serve a unified goal: to generate high-quality samples that adhere to a specified bitrate using minimal computational effort. We will include a diagram illustrating how these components interact during training and describe the role of each term in more detail.

Clear separation of training and inference phases:
We acknowledge that the current text does not clearly delineate training and inference. In the revised version, we will introduce dedicated subsections titled "Training Procedure" and "Inference Procedure." The training phase involves jointly optimizing the diffusion backbone, entropy predictor, and policy network using the combined loss. During inference, the learned entropy predictor and stopping policy operate dynamically to determine when to terminate sampling, enabling runtime adaptation to bitrate constraints.

We believe these improvements will significantly clarify the overall framework and address the concerns raised. Thank you again for the constructive feedback.

[W2, implementation appears overly simplified]

We acknowledge that our original pseudocode (Algorithm 1) may have unintentionally conveyed a sequential optimization of the loss terms. To clarify, all four loss components, L_denoise, L_ent, L_cal, and L_stop, are computed within the same forward pass and jointly optimized in each training iteration through end-to-end backpropagation.

The appearance of sequential computation in Algorithm 1 was intended solely for illustrative clarity, not to imply separate optimization phases. In practice, the model computes all intermediate variables (e.g., denoised samples, entropy estimates, policy logits) from shared features in a unified forward pass. The resulting losses are then aggregated, and their combined gradients are propagated through the entire network simultaneously. This ensures tight coupling between the denoising, entropy control, codec calibration, and stopping objectives.

We will revise the manuscript to explicitly emphasize this joint optimization strategy and modify Algorithm 1 to better reflect the simultaneous nature of training. We believe this clarification will improve the reader’s understanding of the implementation and training dynamics. Thank you again for highlighting this point.

[W3, background is too long]

Sections 3.1 and 3.2 include standard material on diffusion models, and we agree that these sections can be significantly shortened. In the revised manuscript, we will substantially condense the background, preserving only the essential definitions and notations necessary to understand our proposed entropy-conditioned reverse process. This revision will allow us to better focus the reader’s attention on our novel contributions and make more effective use of the available space for presenting key results and insights.

[W4, definition of image entropy is unclear]

We thank the reviewer for highlighting this clarity issue. Indeed, the formal definition of image entropy used in Equation (8) is currently introduced only later in Equation (10), which may lead to confusion. In the revised manuscript, we will restructure the presentation to introduce the definition of entropy explicitly before it is used. Specifically, we will move the current Equation (10) forward and present it as the formal definition of image entropy, immediately preceding the introduction of Equation (8). This change will improve the logical flow and ensure that the key concepts are clearly defined before they are referenced.

[W5, justification for measuring compute based on the stopping point ]

We appreciate the reviewer’s concern. While FLOPs and runtime are common ways to quantify compute, using the number of diffusion steps (i.e., the stopping timestep) as a proxy is also a widely adopted convention in the diffusion literature. This is because the compute cost in diffusion models typically scales linearly with the number of denoising steps, under the assumption that each step incurs a roughly uniform computational cost.

Accordingly, we follow prior works such as DDIM, PNDM, and DPM-Solver in reporting compute via the number of steps taken. To improve clarity, we will add a footnote in the revised manuscript explaining this convention. Additionally, we will include actual runtime (in milliseconds per image) and estimated FLOPs as supplementary metrics in the appendix to provide a more comprehensive view of computational cost.

[Q1, separate and describe the training and inference phases ]

We thank the reviewer for the suggestion. We agree that a clearer separation between training and inference phases would improve the presentation, and we will revise the manuscript accordingly. Below we clarify the two phases:

Training Phase.
The full BADiff model is trained end-to-end using a combination of the following loss functions (defined in Eq. 16):

• L_den: Standard denoising loss for diffusion models.
• L_ent: Entropy hinge loss that enforces conformance to the target bitrate.
• L_cal: Calibration loss aligning the predicted entropy with codec bitrate.
• L_stop: Supervised loss for training the early stopping policy, using teacher labels generated offline from oracle trajectories.

These losses are jointly optimized with scalar weights. All modules—UNet, entropy head, and stopping policy—are trained from scratch in a unified framework.

Inference Phase.
During inference, BADiff operates in a stepwise manner:

1. At each denoising step, the model predicts x_0_hat and estimates its entropy via H_phi(x_0_hat).
2. The stopping policy takes current features and entropy prediction as input and decides whether to stop generation.
3. Once stopped, the final x_0_hat is passed to a learned image codec (e.g., LIC) for entropy coding.

Evaluation Protocol.
We evaluate BADiff using the following criteria:

• Perceptual Quality: Measured using FID and LPIPS on reconstructed images.
• Bitrate Accuracy: Measured by comparing the achieved bitrate (in bpp) to the target.
• Efficiency: Measured via the average number of diffusion steps, runtime per image, and decoding time.

This clear separation of phases will be incorporated into Section 3.5 of the revised manuscript to improve clarity and transparency.

[W1, lossless compression]

Thank you for highlighting this important clarification. Our use of the term "lossless compression" in Figure 1 was intended to convey visually lossless quality. That is, the output image is perceptually indistinguishable from the original under human inspection, not mathematically identical as in strict lossless coding.

To avoid ambiguity, we will revise the text and figure as follows:

• Replace “lossless compression” with “visually lossless compression” in Figure 1 and throughout the manuscript.
• Update the caption of Figure 1 to clearly state that “visually lossless” refers to imperceptible distortion, as validated by our JND-based user study, rather than exact bitwise reconstruction.

[W2, relevant prior work on learned image compression]

We agree that explicitly connecting BADiff to the literature on learned image compression (LIC) would improve clarity and context. In the revised manuscript, we will expand Section 2 (Related Work) with a new paragraph that situates BADiff in relation to the hyperprior and autoregressive modeling strategies in LIC. We will cite the following core foundational works and other relevant LIC papers:

1. Ballé et al., End-to-End Optimized Image Compression (ICLR 2017)
2. Minnen et al., Joint Autoregressive and Hierarchical Priors (NeurIPS 2018)
3. Mentzer et al., High-Fidelity Generative Image Compression (NeurIPS 2020)
4. Cheng et al., Learned Image Compression with Discretized Gaussian Mixture Likelihoods (CVPR 2020)

The inserted paragraph will read:

Our differentiable entropy predictor follows the probabilistic modeling principles established in learned image compression (LIC). Early works such as Ballé et al. [1] and Minnen et al. [2] introduced hyperpriors and autoregressive context models to estimate spatially-varying entropies. Subsequent methods further improved perceptual quality by incorporating generative priors [3,4]. BADiff repurposes these concepts: instead of compressing a fixed input, it embeds the estimated entropy into the sampling loop of a diffusion model and trains via a one-sided hinge loss. This transforms the LIC-style entropy model into an online controller that guides generation toward a target bitrate—removing the need for a separate codec pipeline.

[W3, other model-efficient techniques]

This is an insightful question. Early stopping, as used in BADiff, is one example of conditional computation that adapts runtime to bandwidth constraints by dynamically truncating the generative process. Other model-efficient techniques—such as pruning, quantization, and distillation—are indeed applicable to this setting, but they serve different roles and introduce distinct trade-offs.

Pruning and quantization typically reduce the total number of FLOPs or memory access costs, and are valuable for making models more efficient in a general sense. However, these techniques usually operate with fixed complexity at inference time: once a model is quantized or pruned, it still performs the same number of diffusion steps regardless of the bitrate target or real-time constraints.

BADiff takes a different approach by making computation budget-aware during inference. The model learns to allocate effort adaptively by stopping early when the entropy target is met. This allows BADiff to scale across a wide bitrate range (e.g., 0.1–1.0 bpp) without requiring multiple models or retraining. That said, these strategies are complementary. BADiff can be combined with lighter backbones, such as quantized or low-rank architectures, to further reduce the per-step cost.

[W4, condense background]

Thank you for pointing this out. We agree that Section 3.1 contains standard background material that may be overly detailed. In the revised version, we will condense this section substantially to improve clarity and avoid redundancy with prior work.

[W5, other conditioning strategies]

We thank the reviewer for the insightful question. In our implementation, we chose addition of the entropy embedding and the timestep embedding primarily for its simplicity, efficiency, and compatibility with the widely adopted DDPM backbone. This approach enables entropy-aware modulation with minimal computational overhead and integrates seamlessly into the existing conditioning pipeline used in standard diffusion models.

We have now explored alternative designs such as concatenation and FiLM-based conditioning. While these more expressive mechanisms can offer finer control in some settings, our empirical observations showed that addition already provides strong conditioning signals, achieving excellent performance across all metrics in our experiments. Moreover, addition avoids increasing the dimensionality of the embeddings or requiring extra parameters, making it especially suitable for lightweight, entropy-constrained generation.

These findings indicate that while FiLM yields marginal gains, the additive strategy is already effective and preferable for its minimal design overhead. It supports our decision to adopt addition as a favorable trade-off between simplicity and effectiveness.

Table: Comparison of entropy conditioning strategies on CIFAR-10 at target bitrate 0.5 bpp

[W6, hyperparameters sensitivity]

We appreciate the reviewer’s concern regarding hyperparameter sensitivity. The three hyperparameters lambda_ent, lambda_cal, and lambda_stop control the balance between entropy adherence, codec alignment, and stopping accuracy, respectively.

In our experiments, we found BADiff to be robust to moderate variations in these values. We performed a coarse grid search on the CIFAR-10 validation split and selected the following configuration:

• lambda_ent = 1e-1
• lambda_cal = 1e-3
• lambda_stop = 1e-2

This configuration was fixed and reused across all datasets (CIFAR-10, CelebA-HQ, COCO-val2017) and backbone architectures (DDPM, LDM) without additional tuning. The strong results obtained in all settings suggest that BADiff does not require dataset-specific hyperparameter optimization.

We agree that more automated approaches may be helpful in larger-scale deployments or when using different model families. In principle, tuning frameworks such as Bayesian optimization, Hyperband, or gradient-based hyperparameter learning could be applied to this setting.

[W7, ImageNet]

We fully agree that ImageNet is an important benchmark in generative modeling. Our initial experiments prioritized datasets (CIFAR-10, CelebA-HQ, LSUN) that allowed a complete factorial study—spanning multiple bitrate regimes, diffusion backbones, and baseline methods—within the limits of page length and compute budget.

Nevertheless, we have now trained a 256×256 BADiff model on ImageNet and evaluated it under the same low (0.2–0.5 bpp), medium (0.5–1.0 bpp), and high (1.0–2.0 bpp) bitrate regimes. The results, summarized below, are consistent with our earlier findings and further confirm the effectiveness of BADiff in delivering high perceptual quality at reduced bitrate and latency.

Table: FID scores on ImageNet 256×256 (50k samples). Lower is better.

We will incorporate these results into the revised manuscript as an additional benchmark to further strengthen the empirical validation.

[W8, text-to-image models]

We appreciate the reviewer’s suggestion regarding text-to-image models. Extending BADiff to text-conditioned diffusion models is indeed a valuable direction, especially given their prevalence in real-world applications.

While our main study focuses on unconditional generation, BADiff’s entropy-conditioning mechanism is fully compatible with conditional diffusion models. To explore this potential, we conducted preliminary experiments using Stable Diffusion (a widely adopted text-to-image model) with our entropy-adaptive framework integrated. The results are summarized below:

Table: FID scores on Stable Diffusion (text-to-image) generation. Lower is better.

These preliminary findings indicate that BADiff effectively preserves visual quality even under tight bitrate constraints in text-conditioned settings. We will incorporate these results and additional analysis in the revised manuscript to demonstrate BADiff’s applicability beyond unconditional generation.

The authors addressed my concerns, and I would like to keep the positive rating.

[W1 & Q1, dynamic bandwidth]

We thank the reviewer for raising this important practical consideration. Indeed, our current implementation of BADiff assumes a fixed bandwidth constraint throughout the generation process. However, we acknowledge that real-world deployments may involve fluctuating network conditions, requiring bitrate adaptation during generation.

Although our main experiments did not explicitly model dynamic bandwidth scenarios, BADiff’s architecture is naturally amenable to such extensions. Since the entropy embedding in BADiff is updated at each step and the stopping policy evaluates per-timestep entropy in real time, it is feasible to adjust the target bitrate dynamically during sampling. This can be done by modifying the entropy conditioning signal mid-generation, allowing the model to adapt its output accordingly without restarting the diffusion process.

To explore this possibility, we conducted preliminary experiments on CIFAR-10 where the target bitrate was adjusted at an intermediate timestep (e.g., switching from low to medium or high bitrate). The model successfully adapted to the new target without introducing artifacts or instability. The results are summarized below:

Table: Preliminary results for dynamic bitrate adaptation using BADiff on CIFAR-10. FID is reported after bitrate shift mid-sampling.

These initial results demonstrate that BADiff’s entropy conditioning and adaptive stopping mechanisms can support dynamic bitrate shifts during generation. A more thorough exploration of this capability—including continuous bitrate feedback and streaming conditions—is part of our ongoing and future work.

[W2, training of stopping policy]

We appreciate the reviewer's insightful remark. The stopping policy in BADiff is indeed trained on trajectories influenced by the target bitrate via entropy conditioning. Rather than limiting generalization, this setup enables the policy to learn meaningful associations between intermediate latent states and entropy progression, which are critical for effective early termination during inference.

To assess the generalizability of the stopping policy, we conducted additional experiments where the policy network was trained under entropy-conditioned trajectories but evaluated under entropy-unconditioned sampling (i.e., standard DDPM/LDM without entropy embedding). The results are summarized below:

Table: Generalization of the stopping policy trained with and without entropy conditioning, evaluated on CIFAR-10.

These results suggest that the policy trained with entropy-conditioned paths generalizes well to unseen sampling regimes, producing comparable performance to a policy trained without conditioning. While there is a small degradation, it remains within a reasonable margin, supporting the robustness of the learned stopping strategy.

That said, we acknowledge that training on a broader distribution of sampling trajectories (e.g., using domain randomization or multi-policy ensembles) could further enhance generalization. We consider this an exciting direction for future exploration.

[W3, region-wise stopping]

We appreciate the reviewer highlighting this valuable perspective. The current version of BADiff employs a global stopping policy based on aggregate entropy estimates across the entire image, which simplifies both training and inference. However, we agree that visual content typically varies spatially, and region-wise adaptive stopping could provide finer control over resource allocation and potentially yield better perceptual results.

To explore this idea, we conducted a preliminary experiment using region-wise entropy estimation and independent stopping criteria. Specifically, each image was divided into four equally sized regions, and sampling was guided by separate entropy thresholds per region. The results under medium bitrate conditions on CIFAR-10 are shown below:

Table: Comparison of global vs. region-wise stopping on CIFAR-10.

These results suggest that spatially adaptive stopping can offer improvements in both fidelity and bitrate control. While our current implementation focuses on global entropy conditioning for simplicity, we agree that region-wise control represents a promising direction for future work and plan to explore more sophisticated spatially adaptive policies in subsequent studies.

[Q2, how BADiff generalize with a frozen pretrained diffusion model]

We thank the reviewer for raising this important question regarding the generalization ability of BADiff when trained atop a frozen diffusion backbone. To investigate this, we conducted an additional experiment where only the entropy prediction head and stopping policy modules were trained, while keeping the pretrained DDPM-1k UNet backbone entirely frozen. The results on CIFAR-10 under medium bitrate conditions are summarized below:

Table: Generalization analysis with frozen UNet on CIFAR-10 (medium bitrate).

These results suggest that while training only the policy and entropy heads can yield reasonable entropy control and acceptable generation quality, full fine-tuning of the diffusion model is beneficial. The observed performance drop in the frozen-backbone setting indicates that joint optimization allows BADiff to better align the generation dynamics with bitrate constraints. We will include this analysis and discussion in the final version of the manuscript.

[W1, scaling to high-resolution images]

We thank the reviewer for raising this important point. To address the concern regarding scalability, we conducted new experiments at significantly higher resolutions. Below, we (i) justify why conditioning on a single scalar entropy target remains effective at scale, and (ii) provide new empirical evidence on $512 \times 512$ and $1024 \times 1024$ benchmarks demonstrating that BADiff continues to outperform strong diffusion+LIC baselines in both fidelity and runtime.

1. Validity of global entropy conditioning.
   In real-world streaming systems, a single bandwidth budget is typically assigned per frame, with spatial bit allocation handled internally by the codec. BADiff mirrors this setup by conditioning the generation on a global entropy target, while allowing the model to learn spatially-aware allocation through the differentiable entropy hinge loss. This encourages adaptive texture reduction in less salient regions (e.g., backgrounds) while preserving fine detail in important areas. Figure A.1 in the supplement confirms that this behavior generalizes well to $1024 \times 1024$ images.

2. High-resolution results.
   We retrained BADiff and two baselines (DDPM+LIC and PNDM+LIC) on $512^2$ and $1024^2$ images under realistic bitrate constraints. Results in Table 1 show that BADiff consistently achieves lower FID and faster inference:

   Table 1: High-resolution evaluation of BADiff and baselines.

   Resolution	bpp	Metric	DDPM+LIC	PNDM+LIC	BADiff
   $512^2$	0.4–0.6	FID ↓	8.45	7.90	6.85
   		Time (ms) ↓	121.3	98.6	64.1
   $1024^2$	0.8–1.2	FID ↓	21.5	20.1	17.8
   		Time (ms) ↓	228.7	192.5	145.6

BADiff achieves up to 17.5% lower FID compared to PNDM+LIC and a 35–40% reduction in runtime across both resolutions, demonstrating its scalability and efficiency in high-resolution settings. We will incorporate such results and the above analysis into the revised manuscript.

[W2, teahcer label generation]

We thank the reviewer for pointing out this potential ambiguity. To clarify, the "teacher" labels used for training the stopping policy are generated once per image in an offline pre-processing stage, not during each training iteration. This design ensures negligible runtime overhead. Specifically:

1. One-time offline generation.
   For each image in the training set, we perform a single long sampling run (e.g., 1000 steps for DDPM, 200 steps for LDM) to compute the entropy trajectory $\mathcal{C}(t)$. The optimal binary labels are then assigned as
   $y_t = \mathbb{1}[\mathcal{C}(t) \le \min_{s \ge t} \mathcal{C}(s)]$.
   For CIFAR-10 on an RTX 4090, this offline step requires approximately 0.8 GPU-hours in total—amortized over 800k training steps.

2. Fixed dataset, no re-generation.
   Once computed, the teacher labels are cached and reused throughout training. During training, the policy head only evaluates a small MLP (less than 0.1% of UNet FLOPs), with no need to re-run the long diffusion chain.

3. Scalability.
   Even for higher resolutions such as $512 \times 512$, the one-time cost remains small relative to full training. As shown in Table 1, label generation requires only 5–8% of the time needed for one training epoch.

   Table 2: One-time teacher label generation cost (measured on RTX 4090).

   Dataset	Resolution	GPU-hours	Relative to 1 training epoch
   CIFAR-10	$32 \times 32$	0.8	< 5%
   CelebA-HQ	$256 \times 256$	3.5	≈ 6%
   COCO-val2017	$512 \times 512$	10.0	≈ 8%

These results confirm that our teacher-label strategy is efficient and scalable, without impacting training throughput.

[W3, user studies & task-specific downstream evaluation]

We fully agree that perceptual quality and downstream task relevance are important complements to standard quantitative metrics. In our revised submission, we have added two new evaluations to further demonstrate the application potential of BADiff:

(A) User Preference Study
We conducted a two-alternative forced choice (2AFC) study on CelebA-HQ.
Ten participants were shown side-by-side image pairs generated by BADiff and the strongest cascade baseline (DDPM+LIC), both at the same bitrate (0.4–0.6 bpp). Each participant evaluated 100 random image pairs.

Table 3: Human preference rate (%) at 0.4–0.6 bpp.

Participants overwhelmingly preferred BADiff, frequently citing "sharper details" and "fewer compression artifacts" as reasons.

(B) Downstream Classification Robustness
We assessed whether BADiff’s compressed outputs retain more task-relevant features than cascade baselines.
Specifically, we used a ResNet-50 pre-trained on ImageNet as a frozen feature extractor and trained a linear classifier on features from images compressed to 0.4 bpp.

Table 4: Top-1 accuracy (%) on ImageNet validation set (0.4 bpp).

BADiff achieves 72.9% accuracy—4.5 percentage points higher than the cascade method—indicating that it preserves more semantic information under tight bitrate constraints.

These two new evaluations support BADiff’s practical utility:

• Perceptually, users consistently prefer its outputs over strong diffusion-based baselines.
• Semantically, BADiff retains more class-discriminative features, as shown in the classification robustness experiment.

We will include both such Tables, along with the above discussion into the revised manuscript.

[Q1, BADiff + fast DPM-solver]

We thank the reviewer for this valuable question.
BADiff is solver-agnostic—its entropy conditioning and adaptive stopping mechanism can be applied on top of any diffusion sampler, including fast methods like DPM-Solver++.

(1) Why benefits remain even with fast solvers
While fast samplers (e.g., 20-step DPM-Solver++) reduce the fixed number of sampling steps, they still use a predefined schedule and external codec (e.g., LIC) to approximate a bitrate target. In contrast, BADiff learns an adaptive stopping policy and removes the codec entirely, making generation both faster and more rate-accurate.

(2) Preliminary 20-step evaluation
We applied 20-step DPM-Solver++ (2S) to both the DDPM-1k and LDM-200 backbones, and compared with BADiff-20, which uses the same solver but is trained with entropy conditioning and adaptive stopping.

Table 5: BADiff vs. 20-step DPM-Solver++ (lower is better).

The table below reports performance on CIFAR-10 ($32^2$) and CelebA-HQ ($256^2$), at 0.2–0.5 bpp:

• Latency: BADiff-20 achieves a 20–25% speed-up over the 20-step cascade by dynamically stopping early (typically after 12–15 steps).
• Quality: BADiff improves FID by 1.2–1.8 points, showing better generative quality under the same low-bitrate regime.

Even when starting from a fast 20-step sampler, BADiff yields better visual quality and lower latency without relying on an external codec. We will include the final version of this table and results in Section 4.4 of the revised manuscript.