RISE: Relay Inference and Online Scheduling for Efficient Edge-Device Collaborative Diffusion Model Services

Abstract: Text-to-image diffusion models are increasingly deployed at the network edge to serve heterogeneous workloads with diverse quality and latency requirements. However, existing deployment strategies choose either large edge-side models with high fidelity but high latency or lightweight device-side models that offer speed at the cost of semantic coherence. Moreover, these approaches rarely split the denoising workload between models of different sizes across edge servers and user devices. To bridge this gap, we propose RISE, a method for edge-device diffusion model services that combines relay inference with online scheduling. Driven by the finding that the latent intensity exhibits minimal deviation after a model handoff, RISE uses a training-free relay mechanism that exploits the shared latent space within a model family: the large model on the edge handles the early denoising steps that shape semantic structure, then passes the intermediate latent to a small device-side model for detail refinement. To deploy this mechanism as a practical service, a contextual bandit scheduler selects the best relay configuration based on prompt complexity, user preferences, network quality and real-time node loads. Experiments on two benchmarks show that RISE's relay mechanism achieves up to 2.1$\times$ speedup while preserving full-model quality, and its context-aware scheduler effectively balances quality and latency under mixed workloads.

• The paper presents a novel relay inference approach that splits diffusion model denoising between edge and device, achieving a 2.1× speedup without quality loss.
• The paper introduces a LinUCB-based online scheduler that dynamically selects relay configurations based on prompt complexity, device load, and network conditions.
• The paper validates its method on diffusion model families like SDXL and SD3.5, demonstrating maintained output coherence and highlighting challenges in text rendering in specific setups.

RISE: Relay Inference and Online Scheduling for Efficient Edge-Device Collaborative Diffusion Model Services

Introduction

RISE introduces an edge-device collaborative architecture tailored to diffusion-based text-to-image generation. The method addresses two critical limitations in prior deployments: the inefficiency of running large diffusion models solely at the edge and the fidelity issues encountered with lightweight device-side models. Unlike previous distributed or relay approaches restricted to homogeneous hardware or fixed resolution, RISE leverages the latent space compatibility intrinsic to diffusion model families—such as those in SDXL and Stable Diffusion 3.5 (SD3.5)—to facilitate computation-efficient handoffs. A large model processes the initial, semantics-critical denoising steps; the resulting latent is then relayed to a smaller device model for detail synthesis, all without the necessity for retraining or adaptation.

The framework further introduces a contextual bandit-based (LinUCB) scheduler, which dynamically selects optimal relay configurations (model pair and handoff step) based on prompt complexity, quality-latency tradeoffs, device and edge load, and network state. The primary claim is that this relay and scheduling mechanism achieves a 2.1× speedup without measurable loss in output quality, outperforming both traditional acceleration techniques and alternative collaborative inference baselines.

Empirical Basis for Relay Inference

A fundamental insight driving the relay approach is the observation that early denoising steps in diffusion models establish the core semantic structure, while later steps offer only minor perceptual refinements. This was empirically validated via latent trajectory analysis on the SD3.5 family, where the edge (large) model runs the initial denoising steps, and the device (small) model resumes from a mid-latent state for completion. Quantitative measurement of post-relay latent intensity and per-step deviation demonstrates minimal divergence with handoff, provided the latent spaces are strictly compatible.

Figure 1: Latent intensity comparison between large-model-only and relay configurations reveals close tracking post-handoff, confirming semantic continuity and the feasibility of handoff without retraining.

To operationalize the handoff, RISE implements sigma matching. Given a denoising schedule, it aligns the noise parameter at the handoff so that the device model’s starting step most closely matches the edge model’s current state. This approach handles both identical and asymmetric schedule/mapping cases with minimal error injection.

Figure 2: The process of relay inference divides denoising between edge and device, incorporating sigma matching for seamless latent-state transition.

Relay Mechanism and Model Family Considerations

RISE instantiates the relay mechanism for two canonical diffusion model families: those following the SDXL (UNet backbone, DDIM sampler) and SD3.5 (MMDiT, rectified flow) designs. In SDXL, the edge runs 50 steps and the device runs 25, necessitating nonuniform sigma schedule alignment. In SD3.5, step counts and schedules are identical; thus, the handoff is trivially direct (step-aligned). For both, RISE discretizes candidate relay points (e.g., every five steps), reducing the scheduling action space while capturing dominant quality-latency regimes.

Visual analysis confirms that for general prompts, both quality and coherence are retained post-relay. However, performance for tasks requiring text rendering is strictly model family-dependent, with SD3.5 preserving OCR fidelity through relay and SDXL proving fundamentally incapable across all steps.

Figure 3: Visual comparison across relay configurations illustrating semantic preservation and the importance of model family selection for specific task demands.

Online Context-Aware Relay Scheduling

The second principal component is the online scheduler, which addresses the NP-hard configuration selection in the face of dynamic edge and device conditions and heterogeneous request properties. RISE formulates this as a LinUCB contextual bandit problem, leveraging a compact eight-dimensional context: prompt complexity, text rendering requirements, network quality, battery level, user preference, and resource occupancy. The scheduler maintains per-arm confidence bounds and selects actions via softmax-weighted UCB scoring, naturally balancing exploration and exploitation and adapting over time as task demands and system load evolve.

A dynamic, task-dependent reward integrates multiple objective metrics: visual alignment (CLIP Score), human preference alignment, aesthetic score, OCR accuracy (for text tasks), service latency, VRAM consumption, and device pool congestion. The weights for these metrics are modulated by task context and user preference, ensuring the scheduler aligns with both subjective and service-level priorities.

Experimental Results

The effectiveness of RISE is rigorously validated across two benchmarks (DiffusionDB and DrawTextCreative) and five quality metrics. Comparative evaluation includes both classical and SOTA training-free diffusion acceleration methods (DeepCache, T-GATE, SADA) and several dispatch/scheduling baselines (Round-Robin, Greedy, PPO, SAC).

Numerical Results and Claims

• RISE relay can deliver a speedup up to 2.1× over edge-only inference, while preserving or exceeding full-model quality in CLIP, ImageReward, and PickScore metrics.
• For text-embedded prompts, OCR accuracy is maintained post-relay in SD3.5, but not in SDXL, necessitating context-aware scheduler mediation.
• The online LinUCB-based scheduler outperforms reinforcement learning and heuristic routing by 15.74% in cumulative user-centric reward.
• Ablations reveal context feature and dynamic reward shaping are critical to scheduler efficacy, with total reward dropping up to 49.7% without them.

Figure 4: CLIP Score across configurations, illustrating that relay preserves text-image alignment when model family selection and relay point are properly matched.

Figure 5: Total reward comparison demonstrates the clear performance edge of context-driven LinUCB scheduling against RL baselines and heuristics.

Theoretical and Practical Implications

The primary theoretical contribution is in establishing that denoising-split collaborative inference is feasible and effective without retraining, given strict latent-space compatibility and proper schedule alignment. Practically, RISE enables scalable deployment of high-quality text-to-image services in heterogeneous, bandwidth- and resource-constrained edge environments. The scheduler’s lightweight, strongly performing bandit approach offers a solution that scales, is always-on in a real deployment, and can be robustly adapted to new model families as long as latent compatibility holds.

Notably, the current limitation is the requirement of latent space compatibility, a property generally present within but not across diffusion model families. The authors suggest development of lightweight adapters to bridge cross-family relay, an unsolved but potentially impactful research direction.

Conclusion

RISE establishes a new standard for edge-device collaborative serving of diffusion-based AIGC services by demonstrating that relay inference, coordinated by a context-sensitive bandit scheduler, can simultaneously achieve accelerated inference and quality retention. These results are robust across benchmarks and under real-world, multi-constraint load. Future research can address interoperability between model families and more sophisticated context embeddings for even finer-grade scheduling.