oxicuda_quantum/vqe/layerwise.rs
1//! Layer-wise warm-start initialization to mitigate barren plateaus.
2//!
3//! Reference: Grant, Wossnig, Ostaszewski & Benedetti, "An initialization
4//! strategy for addressing barren plateaus in parametrized quantum circuits",
5//! Quantum 3, 214 (2019); see also Skolik et al. 2021 on layer-wise learning.
6//!
7//! Deep hardware-efficient ansätze suffer from **barren plateaus**: the variance
8//! of the cost gradient vanishes exponentially in the number of qubits, so a
9//! randomly initialized deep circuit is essentially un-trainable. A practical
10//! remedy is to **grow the circuit one layer at a time**:
11//!
12//! 1. Start from a shallow (depth-0) ansatz and optimize it to (near) convergence.
13//! 2. Append one fresh layer whose rotation parameters are **zero-initialized**.
14//! Because `Ry(0) = I`, the appended layer is the identity at initialization,
15//! so the enlarged circuit reproduces the previously optimized state exactly —
16//! the optimizer therefore starts each stage already in a low-cost,
17//! high-gradient region rather than on a random plateau.
18//! 3. Optimize the enlarged parameter set, then repeat until the target depth is
19//! reached.
20//!
21//! The optimizer below realizes this schedule on top of the existing
22//! [`HardwareEfficientAnsatz`] and the parameter-shift energy/gradient already
23//! provided by [`crate::vqe::vqe::VqeOptimizer`].
24
25use crate::error::{QuantumError, QuantumResult};
26use crate::handle::LcgRng;
27use crate::pauli::hamiltonian::Hamiltonian;
28use crate::vqe::ansatz::HardwareEfficientAnsatz;
29use crate::vqe::vqe::VqeOptimizer;
30
31/// Configuration for layer-wise warm-start VQE training.
32#[derive(Debug, Clone)]
33pub struct LayerwiseConfig {
34 /// Number of qubits in the ansatz.
35 pub n_qubits: usize,
36 /// Final ansatz depth to grow to (number of entangling layers).
37 pub target_depth: usize,
38 /// Parameter-shift gradient-descent iterations spent at **each** depth stage.
39 pub iters_per_stage: usize,
40 /// Learning rate for gradient descent.
41 pub lr: f32,
42 /// Small std-dev for randomizing the very first (depth-0) layer; the grown
43 /// layers are always zero-initialized to preserve the identity-block warm
44 /// start, so only the seed layer needs symmetry breaking.
45 pub seed_noise: f32,
46}
47
48impl LayerwiseConfig {
49 /// Validate the configuration.
50 ///
51 /// # Errors
52 /// Returns [`QuantumError::InvalidParameter`] / [`QuantumError::InvalidQubitCount`]
53 /// for non-positive qubit counts, zero iterations, or a non-finite learning rate.
54 pub fn new(
55 n_qubits: usize,
56 target_depth: usize,
57 iters_per_stage: usize,
58 lr: f32,
59 seed_noise: f32,
60 ) -> QuantumResult<Self> {
61 if n_qubits == 0 {
62 return Err(QuantumError::InvalidQubitCount { n: n_qubits });
63 }
64 if iters_per_stage == 0 {
65 return Err(QuantumError::InvalidParameter {
66 name: "iters_per_stage".into(),
67 });
68 }
69 if !lr.is_finite() || lr <= 0.0 {
70 return Err(QuantumError::InvalidParameter { name: "lr".into() });
71 }
72 Ok(Self {
73 n_qubits,
74 target_depth,
75 iters_per_stage,
76 lr,
77 seed_noise,
78 })
79 }
80}
81
82/// Result of a layer-wise warm-start training run.
83#[derive(Debug, Clone)]
84pub struct LayerwiseResult {
85 /// Final optimized parameters for the full target-depth ansatz.
86 pub params: Vec<f32>,
87 /// Final energy ⟨ψ(params)|H|ψ(params)⟩.
88 pub energy: f32,
89 /// Energy after each depth stage, `energies[d]` = energy at depth `d`.
90 pub stage_energies: Vec<f32>,
91}
92
93/// Layer-wise warm-start VQE trainer.
94#[derive(Debug, Clone)]
95pub struct LayerwiseVqe {
96 cfg: LayerwiseConfig,
97 ham: Hamiltonian,
98}
99
100impl LayerwiseVqe {
101 /// Construct the trainer from a validated config and a Hamiltonian.
102 #[must_use]
103 pub fn new(cfg: LayerwiseConfig, ham: Hamiltonian) -> Self {
104 Self { cfg, ham }
105 }
106
107 /// Run the full layer-wise schedule from depth 0 to `target_depth`.
108 ///
109 /// At each stage `d`, the previous stage's optimized parameters are carried
110 /// over verbatim and the freshly added layer's `n_qubits` rotation angles are
111 /// appended as zeros (identity block), guaranteeing the stage starts at the
112 /// previous optimum's energy. Gradient descent (parameter-shift) then refines
113 /// the enlarged parameter vector.
114 ///
115 /// # Errors
116 /// Propagates any error from energy/gradient evaluation.
117 pub fn train(&self, rng: &mut LcgRng) -> QuantumResult<LayerwiseResult> {
118 let mut stage_energies = Vec::with_capacity(self.cfg.target_depth + 1);
119
120 // --- Stage 0: depth-0 ansatz with a small random seed. ---
121 let ansatz0 = HardwareEfficientAnsatz::new(self.cfg.n_qubits, 0);
122 let mut params: Vec<f32> = (0..ansatz0.n_params())
123 .map(|_| rng.next_normal() * self.cfg.seed_noise)
124 .collect();
125 let mut opt = VqeOptimizer {
126 ansatz: ansatz0,
127 ham: self.ham.clone(),
128 params: params.clone(),
129 };
130 let (e0, p0) = opt.optimize(self.cfg.iters_per_stage, self.cfg.lr)?;
131 params = p0;
132 stage_energies.push(e0);
133
134 // --- Grow one entangling layer at a time. ---
135 for depth in 1..=self.cfg.target_depth {
136 let ansatz = HardwareEfficientAnsatz::new(self.cfg.n_qubits, depth);
137 // Carry over previous params and append a zero (identity) RY layer.
138 let mut grown = params.clone();
139 grown.extend(std::iter::repeat_n(0.0_f32, self.cfg.n_qubits));
140 debug_assert_eq!(grown.len(), ansatz.n_params());
141
142 let mut stage_opt = VqeOptimizer {
143 ansatz,
144 ham: self.ham.clone(),
145 params: grown,
146 };
147 let (e_stage, p_stage) = stage_opt.optimize(self.cfg.iters_per_stage, self.cfg.lr)?;
148 params = p_stage;
149 stage_energies.push(e_stage);
150 }
151
152 let final_energy = *stage_energies.last().unwrap_or(&f32::NAN);
153 Ok(LayerwiseResult {
154 params,
155 energy: final_energy,
156 stage_energies,
157 })
158 }
159}
160
161#[cfg(test)]
162mod tests {
163 use super::*;
164 use crate::pauli::pauli_string::PauliOp;
165
166 #[test]
167 fn config_validation() {
168 assert!(LayerwiseConfig::new(0, 2, 5, 0.1, 0.05).is_err());
169 assert!(LayerwiseConfig::new(2, 2, 0, 0.1, 0.05).is_err());
170 assert!(LayerwiseConfig::new(2, 2, 5, 0.0, 0.05).is_err());
171 assert!(LayerwiseConfig::new(2, 2, 5, 0.1, 0.05).is_ok());
172 }
173
174 #[test]
175 fn identity_block_warm_start_does_not_increase_energy_at_grow_boundary() {
176 // The energy after growing should never be worse than before growing
177 // (because the new layer starts as identity, then is optimized).
178 let cfg = LayerwiseConfig::new(2, 3, 8, 0.15, 0.05).expect("valid cfg");
179 let mut ham = Hamiltonian::new();
180 ham.add_term(1.0, vec![PauliOp::Z, PauliOp::Z]);
181 ham.add_term(-0.5, vec![PauliOp::X, PauliOp::I]);
182 let trainer = LayerwiseVqe::new(cfg, ham);
183 let mut rng = LcgRng::new(2024);
184 let res = trainer.train(&mut rng).expect("training succeeds");
185
186 assert_eq!(res.stage_energies.len(), 4); // depths 0,1,2,3
187 // Monotone non-increasing (within numerical slack): deeper warm-started
188 // ansatz can always match the shallower optimum.
189 for w in res.stage_energies.windows(2) {
190 assert!(
191 w[1] <= w[0] + 5e-2,
192 "stage energy increased: {} → {}",
193 w[0],
194 w[1]
195 );
196 }
197 assert!(res.energy.is_finite());
198 }
199
200 #[test]
201 fn final_params_have_target_depth_length() {
202 let cfg = LayerwiseConfig::new(3, 2, 3, 0.1, 0.05).expect("valid cfg");
203 let mut ham = Hamiltonian::new();
204 ham.add_term(1.0, vec![PauliOp::Z, PauliOp::I, PauliOp::Z]);
205 let trainer = LayerwiseVqe::new(cfg, ham);
206 let mut rng = LcgRng::new(11);
207 let res = trainer.train(&mut rng).expect("train");
208 let expected = HardwareEfficientAnsatz::new(3, 2).n_params();
209 assert_eq!(res.params.len(), expected);
210 }
211
212 #[test]
213 fn ground_state_of_single_z_is_recovered() {
214 // H = Z on 1 qubit, ground energy -1. Layer-wise VQE should approach it.
215 let cfg = LayerwiseConfig::new(1, 2, 40, 0.3, 0.1).expect("valid cfg");
216 let mut ham = Hamiltonian::new();
217 ham.add_term(1.0, vec![PauliOp::Z]);
218 let trainer = LayerwiseVqe::new(cfg, ham);
219 let mut rng = LcgRng::new(5);
220 let res = trainer.train(&mut rng).expect("train");
221 assert!(
222 res.energy < -0.9,
223 "energy={} should approach -1",
224 res.energy
225 );
226 }
227}