Introduction
Interventional sampling aims to generate outcomes from the distribution induced by a specified intervention rather than from the observational distribution. This is particularly useful for high-dimensional outcomes. For example, one may wish to generate images with a given demographic attribute A A A A A A [1 , 2 , 3 ] .
Recent approaches construct interventional samplers by specifying a parameterized generative or density-based model, identifying from observational data a training objective that targets the interventional law, and optimizing this objective over the model parameters [4 , 5 , 6 , 7 , 8 ] . This yields flexible samplers, but requires first learning a generative model in addition to the nuisance functions required for causal adjustment. While recent work has simplified the generative learning problem via using flows from data-adaptive source distributions [8 ] , we show the computational cost remains burdensome.
In this work, we introduce Interventional MMD Flow (IMMD-Flow), a method for sampling from interventional distributions which bypasses the need to first learn an explicit generative model. The central idea is to formulate interventional sampling as a sample-based distributional optimization problem, in which an empirical distribution of particles is directly optimized to match the target distribution. We instantiate this using Maximum Mean Discrepancy (MMD), for which the target distribution enters the objective only through its kernel mean embedding [9 , 10 ] . Thus, once the interventional kernel mean embedding is estimated from observational data, samples can be generated directly by evolving the particles under the corresponding MMD gradient flow [11 ] .
The remaining difficulty is statistical: neither the interventional embedding nor samples from the interventional distribution are observed. Under standard identification assumptions, the embedding can be represented through either the treatment-assignment mechanism or a conditional outcome embedding [12 , 13 ] . However, each representation alone is sensitive to nuisance misspecification. We therefore combine them through a doubly robust, Neyman-orthogonal estimator [14 , 15 , 16 ] . This estimator has a simple closed form, can be computed efficiently using a projected Nyström approximation, and has bias of product order in the two nuisance errors. Since the target-dependent component of the squared MMD is linear in the interventional embedding, this second-order robustness carries directly over to the sampling objective. For smooth kernels, the same estimation error also controls perturbations of the induced transport velocity.
Our main contributions are as follows: (i) we formulate interventional sampling as MMD-based distributional optimization and derive a doubly robust, Neyman-orthogonal fixed target whose estimation error perturbs the transport velocity only at product order; (ii) we develop IMMD-Flow and a projected Nyström implementation that reduces conditional embedding regression costs and enables reuse across supported interventions; (iii) for smooth kernels, we derive finite-particle guarantees that propagate target-estimation error through the flow without requiring the estimated target to be a probability embedding, yielding conditions for convergence to the true interventional law; and (iv) we evaluate sampling accuracy, nuisance robustness, and computational performance on synthetic and structured-outcome benchmarks.
§ 2 defines the sampling problem, § 3 develops IMMD-Flow, § 4 gives the practical procedure, and § 5 provides the statistical and particle-level analysis, and § 6 empirically evaluate our method on several benchmarks. We provide a more detailed discussion of related work in Appendix § 11 .
Problem Statement
We study the problem of generating outcomes under a known soft intervention from observational data. We first identify the interventional outcome law and then formulate sampling from this law as optimization over probability measures.
Interventional target and identification. Let X \mathcal{X} X A \mathcal{A} A Y \mathcal{Y} Y P ( Y ) \mathcal{P}(\mathcal{Y}) P ( Y ) Y \mathcal{Y} Y O = ( X , A , Y ) ∼ P 0 O = (X, A, Y) \sim P_0 O = ( X , A , Y ) ∼ P 0 X X X A A A Y Y Y
P 0 ( d y , d a , d x ) = P Y ∣ A , X ( d y ∣ a , x ) π 0 ( d a ∣ x ) P X ( d x ) (1) P_0(\mathrm{d}y, \mathrm{d}a, \mathrm{d}x) = P_{Y \mid A,X}(\mathrm{d}y \mid a, x)\,\pi_0(\mathrm{d}a \mid x)\,P_X(\mathrm{d}x) \tag*{(1)} P 0 ( d y , d a , d x ) = P Y ∣ A , X ( d y ∣ a , x ) π 0 ( d a ∣ x ) P X ( d x ) (1) where P X P_X P X X X X π 0 ( ⋅ ∣ x ) : = P 0 ( A ∈ ⋅ ∣ X = x ) \pi_0(\cdot\mid x) := P_0(A \in\cdot\mid X = x) π 0 ( ⋅ ∣ x ) := P 0 ( A ∈ ⋅ ∣ X = x ) P Y ∣ A , X P_{Y \mid A,X} P Y ∣ A , X
Let { Y ( a ) : a ∈ A } \{Y(a) : a \in\mathcal{A}\} { Y ( a ) : a ∈ A } [18 ] [17]. A fixed, known soft intervention π \pi π x ↦ π ( ⋅ ∣ x ) x \mapsto\pi(\cdot\mid x) x ↦ π ( ⋅ ∣ x ) X \mathcal{X} X A \mathcal{A} A [19 , 20 ] . Its interventional outcome law P π ∈ P ( Y ) P^\pi\in\mathcal{P}(\mathcal{Y}) P π ∈ P ( Y )
P π ( B ) = E [ ∫ A 1 { Y ( a ) ∈ B } π ( d a ∣ X ) ] , (2) P^\pi(B) = \mathbb{E}\left[\int_{\mathcal{A}} 1\{Y(a) \in B\}\,\pi(\mathrm{d}a \mid X)\right], \tag*{(2)} P π ( B ) = E [ ∫ A 1 { Y ( a ) ∈ B } π ( d a ∣ X ) ] , (2) where B ⊆ Y B \subseteq\mathcal{Y} B ⊆ Y P π P^\pi P π π 0 \pi_0 π 0 π \pi π
Assumption 2.1 (Identification). The following standard conditions hold [18 , 19 , 20 ] :
Consistency: Y = Y ( A ) Y = Y(A) Y = Y ( A )
Conditional exchangeability: { Y ( a ) } a ∈ A ⊥ A ∣ X \{Y(a)\}_{a \in\mathcal{A}} \perp A \mid X { Y ( a ) } a ∈ A ⊥ A ∣ X
Overlap: π ( ⋅ ∣ x ) ≪ π 0 ( ⋅ ∣ x ) \pi(\cdot\mid x) \ll\pi_0(\cdot\mid x) π ( ⋅ ∣ x ) ≪ π 0 ( ⋅ ∣ x ) P X P_X P X x x x
Proposition 2.2 (Identification of the interventional law). Under Assumption 2.1 , P π P^\pi P π P 0 P_0 P 0 B ⊆ Y B \subseteq\mathcal{Y} B ⊆ Y
P π ( B ) = ∫ X ∫ A P Y ∣ A , X ( B ∣ a , x ) π ( d a ∣ x ) P X ( d x ) . (3) P^\pi(B) = \int_{\mathcal{X}}\int_{\mathcal{A}} P_{Y \mid A,X}(B \mid a, x)\,\pi(\mathrm{d}a \mid x)\,P_X(\mathrm{d}x). \tag*{(3)} P π ( B ) = ∫ X ∫ A P Y ∣ A , X ( B ∣ a , x ) π ( d a ∣ x ) P X ( d x ) . (3) Since P Y ∣ A , X P_{Y \mid A,X} P Y ∣ A , X P X ( d x ) π 0 ( d a ∣ x ) P_X(\mathrm{d}x)\pi_0(\mathrm{d}a \mid x) P X ( d x ) π 0 ( d a ∣ x ) (3) is invariant to its choice of version. Hence P π P^\pi P π P 0 P_0 P 0 π \pi π
Interventional sampling as distributional optimization. Sampling from P π ∈ P ( Y ) P^\pi\in\mathcal{P}(\mathcal{Y}) P π ∈ P ( Y ) Y ~ ∈ Y \widetilde{Y} \in\mathcal{Y} Y ∈ Y
Pr ( Y ~ ∈ B ) = P π ( B ) , B ⊆ Y measurable . (4) \Pr(\widetilde{Y} \in B) = P^\pi(B), \qquad B \subseteq\mathcal{Y}\ \text{measurable}. \tag*{(4)} Pr ( Y ∈ B ) = P π ( B ) , B ⊆ Y measurable . (4) Our goal is therefore to generate from the full interventional outcome distribution, rather than to estimate a particular functional of it.
Although Proposition 2.2 identifies P π P^\pi P π π 0 \pi_0 π 0 π \pi π [21 , 22 ] . This can be formalized as optimization over probability measures. The resulting formulation can be called sampling distributional optimization (SDO): choose a functional F π : P ( Y ) → R ∪ { + ∞ } F_\pi: \mathcal{P}(\mathcal{Y}) \to\mathbb{R} \cup\{+\infty\} F π : P ( Y ) → R ∪ { + ∞ }
P π ∈ arg min Q ∈ P ( Y ) F π ( Q ) , (5) P^\pi\in\underset{Q \in\mathcal{P}(\mathcal{Y})}{\arg\min} F_\pi(Q), \tag*{(5)} P π ∈ Q ∈ P ( Y ) arg min F π ( Q ) , (5) and construct dynamics that decrease F π F_\pi F π [23 , 22 ] .
Let Q 0 ∈ P ( Y ) Q_0 \in\mathcal{P}(\mathcal{Y}) Q 0 ∈ P ( Y ) ( Q t ) t ≥ 0 (Q_t)_{t \ge0} ( Q t ) t ≥ 0 F π F_\pi F π Q t ⇒ P π Q_t \Rightarrow P^\pi Q t ⇒ P π t → ∞ t \to\infty t → ∞ Y ~ t ∼ Q t \widetilde{Y}_t \sim Q_t Y t ∼ Q t P π P^\pi P π [22 , 24 , 11 ] .
The choice of F π F_\pi F π P π P^\pi P π P π P^\pi P π
Interventional MMD Flow
The sampling distributional optimization formulation in equation SDO formulation requires a distributional objective that identifies the interventional target and is estimable from observational data. We use maximum mean discrepancy (MMD) [9 , 10 ] , whose target dependence is linear in a kernel mean embedding. This structure permits an orthogonal observed-data representation [15 , 16 ] and, subsequently, explicit transport dynamics [11 , 25 , 26 , 27 ] .
MMD objective for interventional sampling
To define the spatial dynamics used below, we work on Z = R d Z \mathcal{Z} = \mathbb{R}^{d_Z} Z = R d Z E : Y → Z E : \mathcal{Y} \to\mathcal{Z} E : Y → Z Z : = E ( Y ) Z := E(Y) Z := E ( Y ) Q π : = E # P π Q^\pi:= E_\#P^\pi Q π := E # P π E # P π E_\#P^\pi E # P π P π P^\pi P π E E E Y = R d Z \mathcal{Y} = \mathbb{R}^{d_Z} Y = R d Z E = Id E = \operatorname{Id} E = Id E E E Q π Q^\pi Q π P π P^\pi P π
Let k : Z × Z → R k : \mathcal{Z} \times\mathcal{Z} \to\mathbb{R} k : Z × Z → R H \mathcal{H} H φ ( z ) : = k ( z , ⋅ ) \varphi(z) := k(z, \cdot) φ ( z ) := k ( z , ⋅ ) Q ∈ P ( Z ) Q \in\mathcal{P}(\mathcal{Z}) Q ∈ P ( Z ) μ Q : = ∫ Z φ ( z ) Q ( d z ) \mu_Q := \int_{\mathcal{Z}} \varphi(z)\,Q(dz) μ Q := ∫ Z φ ( z ) Q ( d z ) χ ( π ) : = μ Q π \chi(\pi) := \mu_{Q^\pi} χ ( π ) := μ Q π H \mathcal{H} H [12 ] , Zenati et al. [13 ] . We take the sampling objective to be
F π ( Q ) : = 1 2 ∥ μ Q − χ ( π ) ∥ H 2 = 1 2 MMD k 2 ( Q , Q π ) . (6) F_\pi(Q) := \frac{1}{2}\lVert\mu_Q - \chi(\pi) \rVert_{\mathcal{H}}^2 = \frac{1}{2}\operatorname{MMD}_k^2(Q,Q^\pi). \tag*{(6)} F π ( Q ) := 2 1 ∥ μ Q − χ ( π ) ∥ H 2 = 2 1 MMD k 2 ( Q , Q π ) . (6) Characteristicness gives F π ( Q ) = 0 F_\pi(Q) = 0 F π ( Q ) = 0 Q = Q π Q = Q^\pi Q = Q π MMD k \operatorname{MMD}_k MMD k Q π Q^\pi Q π [28 ] .
For observational estimation, the useful feature of Equation (6) is its risk representation. Define ℓ Q ( z ) : = 1 2 ∥ μ Q ∥ H 2 − ⟨ μ Q , φ ( z ) ⟩ H \ell_Q(z) := \frac{1}{2}\lVert\mu_Q\rVert_{\mathcal{H}}^2 - \langle\mu_Q,\varphi(z)\rangle_{\mathcal{H}} ℓ Q ( z ) := 2 1 ∥ μ Q ∥ H 2 − ⟨ μ Q , φ ( z ) ⟩ H R π ( Q ) : = ∫ Z ℓ Q ( z ) Q π ( d z ) R_\pi(Q) := \int_{\mathcal{Z}} \ell_Q(z)\,Q^\pi(dz) R π ( Q ) := ∫ Z ℓ Q ( z ) Q π ( d z )
R π ( Q ) − R π ( Q π ) = F π ( Q ) = 1 2 MMD k 2 ( Q , Q π ) . (7) R_\pi(Q) - R_\pi(Q^\pi) = F_\pi(Q) = \frac{1}{2}\operatorname{MMD}_k^2(Q,Q^\pi). \tag*{(7)} R π ( Q ) − R π ( Q π ) = F π ( Q ) = 2 1 MMD k 2 ( Q , Q π ) . (7) Thus, excess interventional risk is exactly half the squared sampling discrepancy. Moreover, the part of the objective that determines optimization depends linearly on χ ( π ) \chi(\pi) χ ( π ) Q Q Q
Orthogonal representation of the interventional objective
Identification expresses χ ( π ) \chi(\pi) χ ( π ) [16 ] .
Write E 0 \mathbb{E}_0 E 0 P 0 P_0 P 0 P A , X ( d a , d x ) = π 0 ( d a ∣ x ) P X ( d x ) P_{A,X}(da,dx)=\pi_0(da\mid x)P_X(dx) P A , X ( d a , d x ) = π 0 ( d a ∣ x ) P X ( d x )
m 0 ( a , x ) : = E 0 [ φ ( Z ) ∣ A = a , X = x ] , w 0 ( a , x ) : = d π ( ⋅ ∣ x ) d π 0 ( ⋅ ∣ x ) ( a ) . m_0(a,x):=\mathbb{E}_0[\varphi(Z)\mid A=a,X=x],\qquad w_0(a,x):=\frac{d\pi(\cdot\mid x)}{d\pi_0(\cdot\mid x)}(a). m 0 ( a , x ) := E 0 [ φ ( Z ) ∣ A = a , X = x ] , w 0 ( a , x ) := d π 0 ( ⋅ ∣ x ) d π ( ⋅ ∣ x ) ( a ) . Proposition 2.2 yields two representations,
χ ( π ) = E 0 [ ∫ A m 0 ( a , X ) π ( d a ∣ X ) ] = E 0 [ w 0 ( A , X ) φ ( Z ) ] . (8) \chi(\pi)=\mathbb{E}_0\left[\int_{\mathcal{A}}m_0(a,X)\pi(da\mid X)\right]=\mathbb{E}_0[w_0(A,X)\varphi(Z)]. \tag*{(8)} χ ( π ) = E 0 [ ∫ A m 0 ( a , X ) π ( d a ∣ X ) ] = E 0 [ w 0 ( A , X ) φ ( Z )] . (8) Either representation in Equation (8) could be used to construct the target embedding. A regression plug-in relies entirely on m 0 m_0 m 0 w 0 w_0 w 0
Instead, we combine the two through an augmented score, following the doubly robust RKHS construction [29 , 13 ] . For candidate nuisance functions m : A × X → H m:\mathcal{A}\times\mathcal{X}\to\mathcal{H} m : A × X → H w : A × X → R w:\mathcal{A}\times\mathcal{X}\to\mathbb{R} w : A × X → R η = ( m , w ) \eta=(m,w) η = ( m , w )
Γ η ( O ) : = ∫ A m ( a , X ) π ( d a ∣ X ) + w ( A , X ) { φ ( Z ) − m ( A , X ) } , (9) \Gamma_\eta(O):=\int_{\mathcal{A}}m(a,X)\pi(da\mid X)+w(A,X)\{\varphi(Z)-m(A,X)\}, \tag*{(9)} Γ η ( O ) := ∫ A m ( a , X ) π ( d a ∣ X ) + w ( A , X ) { φ ( Z ) − m ( A , X )} , (9) with population target h η : = E 0 [ Γ η ( O ) ] ∈ H h_\eta:=\mathbb{E}_0[\Gamma_\eta(O)]\in\mathcal{H} h η := E 0 [ Γ η ( O )] ∈ H
Assumption 3.1 (Nuisance integrability). Let P A , X ( d a , d x ) : = π 0 ( d a ∣ x ) P X ( d x ) P_{A,X}(da,dx):=\pi_0(da\mid x)P_X(dx) P A , X ( d a , d x ) := π 0 ( d a ∣ x ) P X ( d x ) m , m 0 ∈ L 2 ( P A , X ; H ) m,m_0\in L^2(P_{A,X};\mathcal{H}) m , m 0 ∈ L 2 ( P A , X ; H ) w , w 0 ∈ L 2 ( P A , X ) w,w_0\in L^2(P_{A,X}) w , w 0 ∈ L 2 ( P A , X )
Together with boundedness of k k k (9) is Bochner integrable and h η h_\eta h η L 2 L^2 L 2 w 0 w_0 w 0 Assumption 2.1 and controls the degree of intervention-to-observation reweighting.
Importantly, since the interventional MMD risk depends on the target only through χ ( π ) \chi(\pi) χ ( π ) χ ( π ) \chi(\pi) χ ( π ) h η h_\eta h η
R ( Q ; η ) : = 1 2 ∥ μ Q ∥ H 2 − ⟨ μ Q , h η ⟩ H = 1 2 ∥ μ Q − h η ∥ H 2 − 1 2 ∥ h η ∥ H 2 . (10) R(Q;\eta):=\frac{1}{2}\lVert\mu_Q\rVert_{\mathcal{H}}^2-\langle\mu_Q,h_\eta\rangle_{\mathcal{H}}=\frac{1}{2}\lVert\mu_Q-h_\eta\rVert_{\mathcal{H}}^2-\frac{1}{2}\lVert h_\eta\rVert_{\mathcal{H}}^2. \tag*{(10)} R ( Q ; η ) := 2 1 ∥ μ Q ∥ H 2 − ⟨ μ Q , h η ⟩ H = 2 1 ∥ μ Q − h η ∥ H 2 − 2 1 ∥ h η ∥ H 2 . (10) Thus, the same nuisance correction simultaneously defines the target for every candidate distribution Q Q Q
Write r m : = ∥ m − m 0 ∥ L 2 ( P A , X ; H ) r_m:=\lVert m-m_0\rVert_{L^2(P_{A,X};\mathcal{H})} r m := ∥ m − m 0 ∥ L 2 ( P A , X ; H ) r w : = ∥ w − w 0 ∥ L 2 ( P A , X ) r_w:=\lVert w-w_0\rVert_{L^2(P_{A,X})} r w := ∥ w − w 0 ∥ L 2 ( P A , X )
Proposition 3.2 (Double robustness and orthogonality). Under Assumptions 2.1 and 3.1,
h η − χ ( π ) = − E 0 [ ( w − w 0 ) ( m − m 0 ) ] , h_\eta-\chi(\pi)=-\mathbb{E}_0[(w-w_0)(m-m_0)], h η − χ ( π ) = − E 0 [( w − w 0 ) ( m − m 0 )] , ∥ h η − χ ( π ) ∥ H ≤ r w r m , (11) \lVert h_\eta-\chi(\pi)\rVert_{\mathcal{H}}\le r_w r_m, \tag*{(11)} ∥ h η − χ ( π ) ∥ H ≤ r w r m , (11) where the nuisance functions are evaluated at ( A , X ) (A,X) ( A , X ) h m , w 0 = h m 0 , w = χ ( π ) h_{m,w_0}=h_{m_0,w}=\chi(\pi) h m , w 0 = h m 0 , w = χ ( π ) R ( Q ; m , w 0 ) = R ( Q ; m 0 , w ) = R π ( Q ) R(Q;m,w_0)=R(Q;m_0,w)=R_\pi(Q) R ( Q ; m , w 0 ) = R ( Q ; m 0 , w ) = R π ( Q ) Q Q Q
Moreover, viewing R ( Q ; η ) R(Q;\eta) R ( Q ; η ) u = μ Q u=\mu_Q u = μ Q u u u u − h η u-h_\eta u − h η η ˙ \dot{\eta} η ˙
D η ( u − h η ) ∣ η = η 0 [ η ˙ ] = 0 , u ∈ H , (12) D_\eta(u-h_\eta)\big|_{\eta=\eta_0}[\dot{\eta}]=0,\qquad u\in\mathcal{H}, \tag*{(12)} D η ( u − h η ) η = η 0 [ η ˙ ] = 0 , u ∈ H , (12) where η 0 = ( m 0 , w 0 ) \eta_0=(m_0,w_0) η 0 = ( m 0 , w 0 )
The remainder in Equation DR is bilinear in the two nuisance errors, which yields both double robustness and first-order orthogonality. The proof is given in Appendix C.1.
Equation ORTH is the form of Neyman orthogonality relevant to the sampling problem: nuisance perturbations have no first-order effect on the gradient of the risk being optimized [15 , 16 ] . In fact, the property holds for every candidate embedding u u u
Remark 3.3. This construction is related to doubly robust generative-risk correction [7 ] , but MMD yields an additional simplification: because the target-dependent term is linear in the kernel mean embedding, orthogonalizing the entire family of sampling risks reduces to constructing the single target element h η h_\eta h η The next subsection shows that the same element also determines the force driving the MMD transport.
Interventional MMD dynamics
We now turn the orthogonalized MMD objective into a transport sampler. Because the target enters only through a fixed RKHS element, we derive the dynamics for a generic h ∈ H h \in\mathcal{H} h ∈ H h = χ ( π ) h = \chi(\pi) h = χ ( π ) h = h η h = h_\eta h = h η
F h ( Q ) : = 1 2 ∥ μ Q − h ∥ H 2 . (13) \mathcal{F}_h(Q) := \frac{1}{2}\lVert\mu_Q - h \rVert_{\mathcal{H}}^2. \tag*{(13)} F h ( Q ) := 2 1 ∥ μ Q − h ∥ H 2 . (13) When h = χ ( π ) h = \chi(\pi) h = χ ( π ) h = h η h = h_\eta h = h η Q Q Q
To turn Equation 10 into a spatial flow, we impose the following regularity conditions.
Assumption 3.4 (Flow regularity). Let Z = R d Z \mathcal{Z} = \mathbb{R}^{d_Z} Z = R d Z and let Q 0 ∈ P 2 ( Z ) Q_0 \in\mathcal{P}_2(\mathcal{Z}) Q 0 ∈ P 2 ( Z ) The feature map φ : Z → H \varphi: \mathcal{Z} \to\mathcal{H} φ : Z → H is continuously Fréchet differentiable and, for finite constants L φ , L D L_\varphi, L_D L φ , L D
sup z ∈ Z ∥ D φ ( z ) ∥ o p ≤ L φ , ∥ D φ ( z ) − D φ ( z ′ ) ∥ o p ≤ L D ∥ z − z ′ ∥ 2 . \sup_{z \in\mathcal{Z}} \lVert D\varphi(z) \rVert_{\mathrm{op}} \le L_\varphi,\qquad\lVert D\varphi(z) - D\varphi(z') \rVert_{\mathrm{op}} \le L_D \lVert z - z' \rVert_2. z ∈ Z sup ∥ D φ ( z ) ∥ op ≤ L φ , ∥ D φ ( z ) − D φ ( z ′ ) ∥ op ≤ L D ∥ z − z ′ ∥ 2 . For the exact interventional flow, assume additionally Q π ∈ P 2 ( Z ) Q^\pi\in\mathcal{P}_2(\mathcal{Z}) Q π ∈ P 2 ( Z )
These conditions ensure that the velocity fields below are well defined and Lipschitz, so that the corresponding dynamics admit a well-defined measure-valued evolution in P 2 ( Z ) \mathcal{P}_2(\mathcal{Z}) P 2 ( Z ) [22 , 11 ] . It can also be relaxed to local Lipschitzness [30 ] .
For fixed h h h s h ( z ) : = ⟨ h , φ ( z ) ⟩ H s_h(z) := \langle h, \varphi(z) \rangle_{\mathcal{H}} s h ( z ) := ⟨ h , φ ( z ) ⟩ H f Q , h ( z ) : = ⟨ μ Q − h , φ ( z ) ⟩ H = ∫ Z k ( z , z ′ ) Q ( d z ′ ) − s h ( z ) f_{Q,h}(z) := \langle\mu_Q - h, \varphi(z) \rangle_{\mathcal{H}} = \int_{\mathcal{Z}} k(z,z') Q(\mathrm{d}z') - s_h(z) f Q , h ( z ) := ⟨ μ Q − h , φ ( z ) ⟩ H = ∫ Z k ( z , z ′ ) Q ( d z ′ ) − s h ( z ) Q Q Q s h s_h s h [11 ] , samples move along the negative spatial gradient of the witness,
v h ( Q , z ) : = − ∇ f Q , h ( z ) = − ∫ Z ∇ 1 k ( z , z ′ ) Q ( d z ′ ) + ∇ s h ( z ) , v_h(Q,z) := -\nabla f_{Q,h}(z) = -\int_{\mathcal{Z}} \nabla_1 k(z,z') Q(\mathrm{d}z') + \nabla s_h(z), v h ( Q , z ) := − ∇ f Q , h ( z ) = − ∫ Z ∇ 1 k ( z , z ′ ) Q ( d z ′ ) + ∇ s h ( z ) , where ∇ 1 \nabla_1 ∇ 1
We define the Interventional MMD Flow (IMMD-F) as the law-valued evolution induced by v h v_h v h
∂ t Q t h + ∇ ⋅ ( Q t h v h ( Q t h , ⋅ ) ) = 0 , Q t = 0 h = Q 0 . \partial_t Q_t^h + \nabla\cdot\left(Q_t^h v_h(Q_t^h,\cdot)\right) = 0,\qquad Q_{t=0}^h = Q_0. ∂ t Q t h + ∇ ⋅ ( Q t h v h ( Q t h , ⋅ ) ) = 0 , Q t = 0 h = Q 0 . Under Assumption 3.4, this evolution admits the characteristic representation
Z ˙ t = v h ( Q t h , Z t ) , Z 0 ∼ Q 0 , \dot{Z}_t = v_h(Q_t^h,Z_t),\qquad Z_0 \sim Q_0, Z ˙ t = v h ( Q t h , Z t ) , Z 0 ∼ Q 0 , with Z t ∼ Q t h Z_t \sim Q_t^h Z t ∼ Q t h t ≥ 0 t \ge0 t ≥ 0 Q 0 Q_0 Q 0
The generic target h h h h = χ ( π ) h = \chi(\pi) h = χ ( π ) h = h η h = h_\eta h = h η
Proposition 3.5 (Dissipation and velocity stability). Under Assumption 3.4 , any sufficiently regular solution of equation IMMD-F satisfies
d d t F h ( Q t h ) = − ∫ Z ∥ v h ( Q t h , z ) ∥ 2 2 Q t h ( d z ) ≤ 0. (14) \frac{d}{dt}\mathcal{F}_h(Q_t^h)=-\int_{\mathcal{Z}}\|v_h(Q_t^h,z)\|_2^2Q_t^h(dz)\le0. \tag*{(14)} d t d F h ( Q t h ) = − ∫ Z ∥ v h ( Q t h , z ) ∥ 2 2 Q t h ( d z ) ≤ 0. (14) Moreover, for every Q ∈ P 2 ( Z ) Q\in\mathcal{P}_2(\mathcal{Z}) Q ∈ P 2 ( Z )
sup z ∈ Z ∥ v h η ( Q , z ) − v χ ( π ) ( Q , z ) ∥ 2 ≤ L φ ∥ h η − χ ( π ) ∥ H ≤ L φ r w r m . (15) \sup_{z\in\mathcal{Z}}\|v_{h_\eta}(Q,z)-v_{\chi(\pi)}(Q,z)\|_2\le L_\varphi\|h_\eta-\chi(\pi)\|_{\mathcal{H}}\le L_\varphi r_wr_m. \tag*{(15)} z ∈ Z sup ∥ v h η ( Q , z ) − v χ ( π ) ( Q , z ) ∥ 2 ≤ L φ ∥ h η − χ ( π ) ∥ H ≤ L φ r w r m . (15) Equation (14) shows that IMMD Flow decreases the objective associated with its target h h h [11 ] .
Equation VEL gives the dynamical consequence of orthogonality: nuisance error perturbs the transport velocity only at product order. Controlling how this perturbation propagates to trajectories and the generated distribution is deferred to the convergence analysis.
Sampling Procedure
We estimate the interventional target from observational data, then hold it fixed while evolving particles. Note that no nuisance refitting or additional sampling network is required during generation.
Orthogonal target estimation
Given O 1 , … , O n ∼ i i d P 0 O_1,\ldots,O_n\mathrel{\overset{\mathrm{iid}}{\sim}}P_0 O 1 , … , O n ∼ iid P 0 Z i = E ( Y i ) Z_i=E(Y_i) Z i = E ( Y i ) η ^ = ( m ^ , w ^ ) \widehat{\eta}=(\widehat{m},\widehat{w}) η = ( m , w ) (9) :
h ^ π : = 1 n ∑ i = 1 n Γ ^ η ( O i ) . (16) \widehat{h}_{\pi}:=\frac{1}{n}\sum_{i=1}^{n}\widehat{\Gamma}_{\eta}(O_i). \tag*{(16)} h π := n 1 i = 1 ∑ n Γ η ( O i ) . (16) Our implementation fits and evaluates the nuisances on the same observations; its statistical conditions are stated in the convergence analysis.
Projected Nyström construction. To reduce regression cost, we fit m ^ \widehat{m} m [31 ] , using kernels k X , k A k_X,k_A k X , k A m L m_L m L λ > 0 \lambda>0 λ > 0 J J J π ( ⋅ ∣ X i ) \pi(\cdot\mid X_i) π ( ⋅ ∣ X i )
h ~ π = ∑ i = 1 n ν D R , i φ ( Z i ) , s ~ π ( z ) = ∑ i = 1 n ν D R , i k ( z , Z i ) . (17) \widetilde{h}_{\pi}=\sum_{i=1}^{n}\nu_{\mathrm{DR},i}\varphi(Z_i),\qquad\widetilde{s}_{\pi}(z)=\sum_{i=1}^{n}\nu_{\mathrm{DR},i}k(z,Z_i). \tag*{(17)} h π = i = 1 ∑ n ν DR , i φ ( Z i ) , s π ( z ) = i = 1 ∑ n ν DR , i k ( z , Z i ) . (17) Here h ^ π \widehat{h}_{\pi} h π h ~ π \widetilde{h}_{\pi} h π § 12 .§ 2 gives the coefficient construction; the full-kernel alternative is in Appendix § 12 .§ 1 . The coefficients may be signed and need not sum to one, so zero residual against h ~ π \widetilde{h}_{\pi} h π
Particle generation
We instantiate IMMD Flow with h = h ~ π h=\widetilde{h}_{\pi} h = h π Y j ( 0 ) ∼ i i d P i n i t Y_j^{(0)}\mathrel{\overset{\mathrm{iid}}{\sim}}P_{\mathrm{init}} Y j ( 0 ) ∼ iid P init Z j ( 0 ) = E ( Y j ( 0 ) ) Z_j^{(0)}=E(Y_j^{(0)}) Z j ( 0 ) = E ( Y j ( 0 ) ) P i n i t P_{\mathrm{init}} P init Q 0 = E # P i n i t ∈ P 2 ( Z ) Q_0=E_{\#}P_{\mathrm{init}}\in\mathcal{P}_2(\mathcal{Z}) Q 0 = E # P init ∈ P 2 ( Z ) r r r Q ^ M ( r ) = M − 1 ∑ j = 1 M δ Z j ( r ) \widehat{Q}_M^{(r)}=M^{-1}\sum_{j=1}^{M}\delta_{Z_j^{(r)}} Q M ( r ) = M − 1 ∑ j = 1 M δ Z j ( r )
Following Arbel et al. [11 ] and Chen et al. [27 ] , we optionally perturb the gradient evaluation locations: set Z ~ j ( r ) = Z j ( r ) + β r U j ( r ) \widetilde{Z}_j^{(r)}=Z_j^{(r)}+\beta_rU_j^{(r)} Z j ( r ) = Z j ( r ) + β r U j ( r ) U j ( r ) ∼ N ( 0 , I d Z ) U_j^{(r)}\sim\mathcal{N}(0,I_{d_Z}) U j ( r ) ∼ N ( 0 , I d Z ) β r = c β ( r + 1 ) − α \beta_r=c_\beta(r+1)^{-\alpha} β r = c β ( r + 1 ) − α c β ≥ 0 c_\beta\ge0 c β ≥ 0 α ∈ [ 0 , 1 / 2 ) \alpha\in[0,1/2) α ∈ [ 0 , 1/2 ) η r > 0 \eta_r>0 η r > 0
G j ( r ) = 1 M ∑ ℓ = 1 M ∇ 1 k ( Z ~ j ( r ) , Z ℓ ( r ) ) − ∇ s ~ π ( Z ~ j ( r ) ) , (18) G_j^{(r)}=\frac{1}{M}\sum_{\ell=1}^{M}\nabla_1 k\left(\widetilde{Z}_j^{(r)},Z_\ell^{(r)}\right)-\nabla\widetilde{s}_{\pi}\left(\widetilde{Z}_j^{(r)}\right), \tag*{(18)} G j ( r ) = M 1 ℓ = 1 ∑ M ∇ 1 k ( Z j ( r ) , Z ℓ ( r ) ) − ∇ s π ( Z j ( r ) ) , (18) Z j ( r + 1 ) = Z j ( r ) − η r G j ( r ) . Z_j^{(r+1)}=Z_j^{(r)}-\eta_rG_j^{(r)}. Z j ( r + 1 ) = Z j ( r ) − η r G j ( r ) . Algorithm 1 Projected-Nyström IMMD sampler
Require: Data O 1 : n O_{1:n} O 1 : n π \pi π E E E D D D k X k_X k X k A k_A k A k k k m L m_L m L J J J λ \lambda λ M M M R R R P i n i t P_{\mathrm{init}} P init ( η r , β r ) (\eta_r, \beta_r) ( η r , β r )
Set Z i ← E ( Y i ) Z_i \leftarrow E(Y_i) Z i ← E ( Y i ) i i i
Fit w ^ \widehat{w} w w ^ ← ( w ^ ( A i , X i ) ) i = 1 n \widehat{w} \leftarrow(\widehat{w}(A_i, X_i))_{i=1}^n w ← ( w ( A i , X i ) ) i = 1 n
Select m L m_L m L
Draw J J J
Compute b p r o j b^{\mathrm{proj}} b proj § 12 .§ 2 .
Set v D R ← ( b p r o j + w ^ ) / n v_{\mathrm{DR}} \leftarrow(b^{\mathrm{proj}} + \widehat{w})/n v DR ← ( b proj + w ) / n s ^ π \widehat{s}_{\pi} s π (17) .
Draw Y j ( 0 ) ∼ i i d P i n i t Y_j^{(0)} \mathrel{\overset{\mathrm{iid}}{\sim}} P_{\mathrm{init}} Y j ( 0 ) ∼ iid P init Z j ( 0 ) ← E ( Y j ( 0 ) ) Z_j^{(0)} \leftarrow E(Y_j^{(0)}) Z j ( 0 ) ← E ( Y j ( 0 ) ) j j j
for r = 0 , … , R − 1 r = 0, \ldots, R - 1 r = 0 , … , R − 1 do
Draw independent U j ( r ) ∼ N ( 0 , I d Z ) U_j^{(r)} \sim\mathcal{N}(0, I_{d_Z}) U j ( r ) ∼ N ( 0 , I d Z )
Set Z ~ j ( r ) ← Z j ( r ) + β r U j ( r ) \widetilde{Z}_j^{(r)} \leftarrow Z_j^{(r)} + \beta_r U_j^{(r)} Z j ( r ) ← Z j ( r ) + β r U j ( r ) j j j
Compute all G j ( r ) G_j^{(r)} G j ( r ) (18) .
end for
Set Y ^ j ← D ( Z j ( R ) ) \widehat{Y}_j \leftarrow D(Z_j^{(R)}) Y j ← D ( Z j ( R ) ) j j j
return Q ^ M ( R ) \widehat{Q}_M^{(R)} Q M ( R ) P ^ M \widehat{P}_M P M
All gradients use the same current particle cloud; only their evaluation locations are perturbed. Setting c β = 0 c_\beta= 0 c β = 0
After R R R D : Z → Y D : \mathcal{Z} \to\mathcal{Y} D : Z → Y Q ^ M ( R ) \widehat{Q}_M^{(R)} Q M ( R ) P ^ M = D # Q ^ M ( R ) \widehat{P}_M = D_{\#}\widehat{Q}_M^{(R)} P M = D # Q M ( R ) E = D = I d E = D = \mathrm{Id} E = D = Id § 15 .§ 2 ).
Statistical Guarantees
In this section, we bound the discrepancy between the generated empirical distribution and Q π Q^\pi Q π h ^ π \widehat{h}_\pi h π (16) is the doubly robust target estimator with intervention integrals evaluated exactly, whereas h ~ π \widetilde{h}_\pi h π (17) is its numerical approximation. Write
e n , J : = ∥ h ~ π − χ ( π ) ∥ H . e_{n,J} := \lVert\widetilde{h}_\pi- \chi(\pi) \rVert_{\mathcal{H}}. e n , J := ∥ h π − χ ( π ) ∥ H . We first control sampling error in terms of e n , J e_{n,J} e n , J § 5.1 , then bound this quantity using orthogonal target estimation in § 5.2 . The representation, decoder, and kernels are fixed independently of the observations analyzed here.
IMMD-Flow sampling error
A growing body of work studies the MMD-flow error, establishing conditions under which it vanishes as the number of particles M M M R R R [32 ] , noise injection [11 , 27 ] , adaptive kernel lengthscales [33 ] , Fisher–Rao geometry [34 , 35 ] , and energy/Riesz kernels [36 , 37 , 38 , 39 ] . For finite-particle analysis, we consider the noise-injection framework of Arbel et al. [11 ] , Chen et al. [27 ] , because it gives explicit finite-iteration and finite-particle bounds for smooth kernels.
Assumption 5.1 (Smooth particle dynamics). The feature map φ \varphi φ κ ≥ 1 \kappa\ge1 κ ≥ 1 η r = τ > 0 \eta_r = \tau> 0 η r = τ > 0 β r = c β ( r + 1 ) − α \beta_r = c_\beta(r+1)^{-\alpha} β r = c β ( r + 1 ) − α
12 τ κ 2 d Z ≤ 1 , 0 < c β ≤ 1 2 d Z , 0 ≤ α < 1 2 . 12\tau\kappa^2 d_Z \le1,\qquad0 < c_\beta\le\frac{1}{2\sqrt{d_Z}},\qquad0 \le\alpha< \frac{1}{2}. 12 τ κ 2 d Z ≤ 1 , 0 < c β ≤ 2 d Z 1 , 0 ≤ α < 2 1 . Smoothness alone does not ensure convergence to the target. As in noise-injected MMD-flow analysis [11 , 27 ] , we additionally impose a gradient-dominance condition over the iterations being analyzed. Let D \mathcal{D} D F r \mathcal{F}_r F r r r r
Assumption 5.2 (Noisy gradient dominance). For r = 0 , … , R − 1 r = 0, \ldots, R - 1 r = 0 , … , R − 1 h ~ π \widetilde{h}_{\pi} h π
1 M ∑ j = 1 M E [ ∥ v χ ( π ) ( Q ^ M ( r ) , Z j ( r ) + β r U j ( r ) ) ∥ 2 2 | F r ] ≥ 4 κ 2 d Z β r 2 MMD k 2 ( Q ^ M ( r ) , Q π ) , (19) \frac{1}{M} \sum_{j=1}^{M} \mathbb{E}\left[\left\|v_{\chi(\pi)}\left(\widehat{Q}^{(r)}_{M}, Z^{(r)}_{j} + \beta_r U^{(r)}_{j}\right)\right\|_{2}^{2} \,\middle|\, \mathcal{F}_r\right] \ge4\kappa^{2}d_Z\beta_r^{2}\operatorname{MMD}_{k}^{2}\left(\widehat{Q}^{(r)}_{M}, Q^{\pi}\right), \tag*{(19)} M 1 j = 1 ∑ M E [ v χ ( π ) ( Q M ( r ) , Z j ( r ) + β r U j ( r ) ) 2 2 F r ] ≥ 4 κ 2 d Z β r 2 MMD k 2 ( Q M ( r ) , Q π ) , (19) where U j ( r ) U^{(r)}_{j} U j ( r )
This condition evaluates the true-target velocity along the estimated-target particle process. It is an additional optimization assumption, rather than a consequence of the prescribed noise schedule or nuisance orthogonality. We impose this condition to establish a favorable convergence rate, rather than because it is necessary for the good performance of MMD flows. Indeed, as shown by Logarajah et al. [40 ] , a sufficiently large number of interacting particles can also guarantee good performance without a prescribed noise schedule.
Theorem 5.3 (Estimated-target particle bound). Under Assumptions 5.1 and 5.2 , there exist constants C 0 , C 1 , C 2 > 0 C_0,C_1,C_2 > 0 C 0 , C 1 , C 2 > 0 κ , d Z , τ , c β , α \kappa,d_Z,\tau,c_\beta,\alpha κ , d Z , τ , c β , α R ≥ 1 R \ge1 R ≥ 1 δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 )
MMD k ( Q ^ M ( R ) , Q π ) ≤ C 0 exp ( − τ κ 2 d Z c β 2 16 ( 1 − 2 α ) R 1 − 2 α ) + ( 1 + C 1 R α ) e n , J + C 2 R α log ( R / δ ) M . (20) \operatorname{MMD}_{k}\left(\widehat{Q}^{(R)}_{M},Q^{\pi}\right) \le C_0 \exp\left(-\frac{\tau\kappa^{2}d_Zc_\beta^{2}}{16(1-2\alpha)}R^{1-2\alpha}\right) + (1+C_1R^\alpha)e_{n,J} + C_2R^\alpha\sqrt{\frac{\log(R/\delta)}{M}}. \tag*{(20)} MMD k ( Q M ( R ) , Q π ) ≤ C 0 exp ( − 16 ( 1 − 2 α ) τ κ 2 d Z c β 2 R 1 − 2 α ) + ( 1 + C 1 R α ) e n , J + C 2 R α M log ( R / δ ) . (20) For each realization of D \mathcal{D} D ∥ h ~ π ∥ H ≤ κ \|\widetilde{h}_{\pi}\|_{\mathcal{H}} \le\kappa ∥ h π ∥ H ≤ κ 1 − δ 1-\delta 1 − δ
Proof is given in Appendix D.2. The three terms describe optimization decay, propagated numerical target error, and finite-particle error. Increasing the iteration number R R R ∥ h ~ π ∥ H ≤ κ \|\widetilde{h}_{\pi}\|_{\mathcal{H}} \le\kappa ∥ h π ∥ H ≤ κ (125) .
Statistical and numerical target error
To apply Theorem 5.3 , we control the numerical target relative to χ ( π ) \chi(\pi) χ ( π ) η ^ n = ( m ^ n , w ^ n ) \widehat{\eta}_n=(\widehat{m}_n,\widehat{w}_n) η n = ( m n , w n ) P n g = n − 1 ∑ i g ( O i ) P_n g=n^{-1}\sum_i g(O_i) P n g = n − 1 ∑ i g ( O i ) P 0 g = E 0 [ g ( O ) ] P_0g=\mathbb{E}_0[g(O)] P 0 g = E 0 [ g ( O )]
Assumption 5.4 (Nuisance estimation and sample reuse). For deterministic rates r m , n , r w , n r_{m,n},r_{w,n} r m , n , r w , n
∥ m ^ n − m 0 ∥ L 2 ( P A , X ; H ) = O p ( r m , n ) , ∥ w ^ n − w 0 ∥ L 2 ( P A , X ) = O p ( r w , n ) . \|\widehat{m}_n-m_0\|_{L^2(P_{A,X};\mathcal{H})}=O_p(r_{m,n}),\qquad\|\widehat{w}_n-w_0\|_{L^2(P_{A,X})}=O_p(r_{w,n}). ∥ m n − m 0 ∥ L 2 ( P A , X ; H ) = O p ( r m , n ) , ∥ w n − w 0 ∥ L 2 ( P A , X ) = O p ( r w , n ) . The oracle score has finite second moment, and
∥ ( P n − P 0 ) { Γ η ^ n − Γ η 0 } ∥ H = O p ( n − 1 / 2 ) , (21) \|(P_n-P_0)\{\Gamma_{\widehat{\eta}_n}-\Gamma_{\eta_0}\}\|_{\mathcal{H}}=O_p(n^{-1/2}), \tag*{(21)} ∥ ( P n − P 0 ) { Γ η n − Γ η 0 } ∥ H = O p ( n − 1/2 ) , (21) where η 0 = ( m 0 , w 0 ) \eta_0=(m_0,w_0) η 0 = ( m 0 , w 0 )
Condition Equation (21) accounts for fitting and evaluating the nuisances on the same observations; a regression prediction rate alone does not establish it. For intervention integration, define m ‾ n ( X i ) : = ∫ m ^ n ( a , X i ) π ( d a ∣ X i ) \overline{m}_n(X_i):=\int\widehat{m}_n(a,X_i)\,\pi(da\mid X_i) m n ( X i ) := ∫ m n ( a , X i ) π ( d a ∣ X i )
V n : = 1 n ∑ i = 1 n ∫ ∥ m ^ n ( a , X i ) − m ‾ n ( X i ) ∥ H 2 π ( d a ∣ X i ) . V_n:=\frac{1}{n}\sum_{i=1}^{n}\int\|\widehat{m}_n(a,X_i)-\overline{m}_n(X_i)\|_{\mathcal{H}}^{2}\,\pi(da\mid X_i). V n := n 1 i = 1 ∑ n ∫ ∥ m n ( a , X i ) − m n ( X i ) ∥ H 2 π ( d a ∣ X i ) . Proposition 5.5 (Accuracy of the numerical target). Under Assumptions 2.1 , 3.1, and 5.4 , suppose that, conditional on the observations, landmarks, and fitted nuisances, intervention integration uses J J J V n = O p ( 1 ) V_n=O_p(1) V n = O p ( 1 ) J J J
e n , J = O p ( n − 1 / 2 + r m , n r w , n + ( n J ) − 1 / 2 ) . e_{n,J}=O_p\left(n^{-1/2}+r_{m,n}r_{w,n}+(nJ)^{-1/2}\right). e n , J = O p ( n − 1/2 + r m , n r w , n + ( n J ) − 1/2 ) . The product term follows from equation DR; the Monte Carlo term follows from the conditional mean-square error V n / ( n J ) V_n/(nJ) V n / ( n J ) J ≥ 1 J \ge1 J ≥ 1 V n V_n V n Theorem D.16 , substituting r m , n = n − 1 / [ 2 ( 1 + γ ) ] r_{m,n}=n^{-1/[2(1+\gamma)]} r m , n = n − 1/ [ 2 ( 1 + γ )]
e n , J = O p ( n − 1 / 2 + r w , n n − 1 2 ( 1 + γ ) + ( n J ) − 1 / 2 ) , (22) e_{n,J}=O_p\left(n^{-1/2}+r_{w,n}n^{-\frac{1}{2(1+\gamma)}}+(nJ)^{-1/2}\right), \tag*{(22)} e n , J = O p ( n − 1/2 + r w , n n − 2 ( 1 + γ ) 1 + ( n J ) − 1/2 ) , (22) where γ \gamma γ Assumption D.5 . This specialization requires the stated ridge and growing-landmark conditions.
Consistency
Now, we are ready to combine Theorem 5.3 and Proposition 5.5 to obtain the final consistency result. To give an explicit choice of computational budgets R R R M M M n n n r m , n r w , n = O ( n − s ) r_{m,n}r_{w,n}=O(n^{-s}) r m , n r w , n = O ( n − s ) s ∈ ( 0 , 1 / 2 ] s\in(0,1/2] s ∈ ( 0 , 1/2 ]
Corollary 5.6. Suppose the assumptions of Theorem 5.3 and Proposition 5.5 hold over the chosen iteration horizons, with constants independent of n n n κ > sup z ∈ Z ∥ φ ( z ) ∥ H \kappa>\sup_{z\in\mathcal{Z}}\|\varphi(z)\|_{\mathcal{H}} κ > sup z ∈ Z ∥ φ ( z ) ∥ H r m , n r w , n = O ( n − s ) r_{m,n}r_{w,n}=O(n^{-s}) r m , n r w , n = O ( n − s ) s ∈ ( 0 , 1 / 2 ] s\in(0,1/2] s ∈ ( 0 , 1/2 ] J ≥ 1 J\ge1 J ≥ 1 C R = 16 s ( 1 − 2 α ) τ κ 2 d Z c β 2 C_R=\frac{16s(1-2\alpha)}{\tau\kappa^2d_Zc_\beta^2} C R = τ κ 2 d Z c β 2 16 s ( 1 − 2 α )
R n = ⌈ ( C R log n ) 1 / ( 1 − 2 α ) ⌉ , M n = ⌈ n 2 s log log n ⌉ , (23) R_n=\left\lceil(C_R\log n)^{1/(1-2\alpha)}\right\rceil,\qquad M_n=\left\lceil n^{2s}\log\log n\right\rceil, \tag*{(23)} R n = ⌈ ( C R log n ) 1/ ( 1 − 2 α ) ⌉ , M n = ⌈ n 2 s log log n ⌉ , (23) Then, MMD k ( Q ^ M n ( R n ) , Q π ) = O p ( n − s ( log n ) α / ( 1 − 2 α ) ) \operatorname{MMD}_k(\widehat{Q}_{M_n}^{(R_n)},Q^\pi)=O_p\left(n^{-s}(\log n)^{\alpha/(1-2\alpha)}\right) MMD k ( Q M n ( R n ) , Q π ) = O p ( n − s ( log n ) α / ( 1 − 2 α ) )
If MMD k \operatorname{MMD}_k MMD k [28 ] , the generated empirical distribution also converges weakly to the interventional distribution Q π Q^\pi Q π
Finally, for encoded outcomes, the remaining step is to transfer latent accuracy through the decoder.
Proposition 5.7 (Outcome-space transfer). Let k E ( y , y ′ ) : = k ( E ( y ) , E ( y ′ ) ) k_E(y,y'):=k(E(y),E(y')) k E ( y , y ′ ) := k ( E ( y ) , E ( y ′ )) E , D E,D E , D φ \varphi φ L k L_k L k MMD k E ( D # Q , P π ) ≤ MMD k ( Q , Q π ) + L k E Z ∼ Q ∥ E ( D ( Z ) ) − Z ∥ 2 \operatorname{MMD}_{k_E}(D_{\#}Q,P^\pi)\le\operatorname{MMD}_k(Q,Q^\pi)+L_k\mathbb{E}_{Z\sim Q}\|E(D(Z))-Z\|_2 MMD k E ( D # Q , P π ) ≤ MMD k ( Q , Q π ) + L k E Z ∼ Q ∥ E ( D ( Z )) − Z ∥ 2
Applying this bound to Q = Q ^ M ( R ) Q=\widehat{Q}_{M}^{(R)} Q = Q M ( R ) E E E k E k_E k E D D D D # Q π = P π D_{\#}Q^\pi=P^\pi D # Q π = P π E = D = Id E=D=\operatorname{Id} E = D = Id
Numerical Experiments
Methods. We evaluate our orthogonal variant of our IMMD-Flow method, that we write DR–IMMD (for the Doubly Robust embedding) together with two variants of MMD flow using plug-in estimates of the interventional embedding. Plug-in–IMMD constructs the fixed target using only the fitted conditional mean embedding, whereas IPW–IMMD uses only the estimated intervention-to-logging density ratio. DR–IMMD combines both through the orthogonal score in § 3 . These variants differ only in the target coefficients supplied to the flow. Their exact constructions are given in Appendix § 12 .
We compare against Deconfounding Flow Matching (DecFM) [8 ] , which is formulated for arm-specific counterfactual distributions corresponding to hard interventions A ← a A\leftarrow a A ← a π ( ⋅ ∣ x ) = δ a \pi(\cdot\mid x)=\delta_a π ( ⋅ ∣ x ) = δ a E E E D D D π ( ⋅ ∣ X ) \pi(\cdot\mid X) π ( ⋅ ∣ X )
Benchmarks. We evaluate the methods on three image-valued deconfounding tasks. ColorMNIST, adapted from Dance et al. [8 ] , couples treatment-specific digit identity with a context variable that affects both treatment assignment and foreground colour. The interventional target fixes the digit class while averaging over the population distribution of the colour-generating context, thereby removing the observational colour imbalance.
The CelebA benchmark of Dance et al. [8 ] sets A = S e x A = \mathrm{Sex} A = Sex X = H a i r C o l o r X = \mathrm{HairColor} X = HairColor a a a P X ∣ A = a P_{X\mid A=a} P X ∣ A = a P Y ∣ A = a P_{Y\mid A=a} P Y ∣ A = a P X P_X P X A = a A=a A = a [13 ] over soft interventions. The context determines sprite shape and scale and also affects a stochastic action controlling position. The target policy changes the conditional action distribution and hence the joint distribution of visual content and position. Full descriptions of the data-generating processes are given in Appendix § 15 .§ 3 .
Results. We evaluate generated samples using empirical sliced 2-Wasserstein distance in decoded pixel space (S W 2 SW_2 S W 2 S W 2 SW_2 S W 2 [41 ] , together with FID on decoded images [42 ] ; lower is better. All experiments in this section use n = 20,000 n = 20{,}000 n = 20 , 000 M = 2,000 M = 2{,}000 M = 2 , 000
Sample quality. Table 1 shows that all IMMD variants outperform every DecFM baseline in pixel-space S W 2 SW_2 S W 2 S W 2 SW_2 S W 2
Table 1. Comparison on ColorMNIST, CelebA, and dSprites. Lower is better; both S W 2 SW_2 S W 2 10 3 10^3 1 0 3 ± \pm ±
Computational efficiency. A central advantage of IMMD-Flow is that it avoids training an intervention-specific generative model. Once the nuisance functions have been estimated, the target kernel embedding is constructed directly and held fixed, and samples are generated by evolving a finite collection of particles under the resulting MMD flow. In contrast, DecFM learns a neural velocity field through flow matching for the interventional distribution.
Table 2 compares one-time fitting and post-fitting sample generation on ColorMNIST and CelebA. For
Table 2. End-to-end wall-clock comparison at n = 20,000 n = 20{,}000 n = 20 , 000 M = 2,000 M = 2{,}000 M = 2 , 000 ± \pm ±
DR–IMMD, fitting comprises learning the nuisance functions—the conditional mean embedding and intervention-to-logging density ratio—and constructing the fixed orthogonal RKHS target. During sampling, these quantities remain fixed while M = 2,000 M = 2{,}000 M = 2 , 000 n = 20,000 n = 20{,}000 n = 20 , 000 M = 2,000 M = 2{,}000 M = 2 , 000
Particle evolution. Following Dance et al. [8 ] , we initialize particles from the observational distribution P Y ∣ A = a P_{Y\mid A=a} P Y ∣ A = a Figure 1 illustrates the resulting deconfounding process on ColorMNIST. As the particles evolve, their foreground colours change while their digit identities are largely preserved. At the population level, the colour distribution moves from the confounded observational law towards the oracle interventional distribution P X P_X P X
Figure 1. ColorMNIST DR–IMMD at n = 20,000 n = 20{,}000 n = 20 , 000 M = 2,000 M = 2{,}000 M = 2 , 000 P Y ∣ A = a P_{Y\mid A=a} P Y ∣ A = a § 15 .§ 5 .
Figure 2 illustrates a similar process on CelebA, where particles undergo gradual changes in appearance, including hair colour, while much of their original facial structure remains recognisable. In both cases, these changes arise from optimizing the particle locations against a fixed estimate of the interventional kernel embedding, with the decoder held fixed throughout. The resulting trajectories illustrate how initializing from observational samples allows the flow to exploit their existing structure when approximating the interventional distribution.
Figure 2. Decoded CelebA particles during DR–IMMD optimization at n = 20,000 n = 20{,}000 n = 20 , 000 M = 2,000 M = 2{,}000 M = 2 , 000 P Y ∣ A = a P_{Y\mid A=a} P Y ∣ A = a P Y ( a ) P_{Y(a)} P Y ( a ) § 5 .
Ablations. Appendices § 15 .§ 5 and § 15 .§ 6 provide additional qualitative samples, particle trajectories, and ablations over sample size, particle count, initialization, confounding, and Nyström approximation. Sampling generally improves with more logged data and consistently with more particles. Observational initialization is particularly beneficial on the harder image benchmarks, while performance degrades only gradually as overlap weakens.
Conclusion
We introduced IMMD-Flow, which turns an estimated interventional RKHS target into samples through particle transport, without training an additional sampler. The framework accommodates both supported hard interventions and context-dependent stochastic policies, including continuous-action policies. Its doubly robust, orthogonal target yields product-order perturbations of the transport dynamics. Projected Nyström enables reuse across interventions. On ColorMNIST and CelebA, IMMD-Flow improves pixel-space S W 2 \mathrm{SW}_2 SW 2 11 11 11 42 × 42\times 42 ×
Code Availability
Our implementation and code for reproducing the experiments are available at https://github.com/BariscanBozkurt/Interventional-MMD-Flow-Sampling .
AI use statement
Generative AI tools were used as editorial and technical aids in preparing this paper. Their use included refining the manuscript’s language and structure, supporting code modifications, generating simulation scripts, generating Bash scripts to launch simulations in SLURM cluster, and drafting and checking mathematical concepts and proofs. The authors examined all AI-assisted outputs and independently validated the final text, implementation, and mathematical arguments. The authors take responsibility for the final content of the paper and the accompanying implementation
References
Discussion of Related work
Debiased estimation of interventional distributions. A growing literature extends semiparametric debiasing methods from scalar causal estimands to entire interventional or counterfactual outcome distributions. Early work developed efficient estimators for projections and functionals of counterfactual densities, including truncated-series and exponential-family approximations [4 ] . More recent methods combine these ideas with flexible generative models. Interventional Normalizing Flows Melnychuk et al. [5 ] parameterize the target interventional density by a normalizing flow and optimize a one-step corrected likelihood objective, while Martinez-Taboada and Kennedy [6 ] use a doubly robust Kernel Stein Discrepancy to estimate an energy-based counterfactual density. Luedtke and Fukumizu [7 ] give a more general doubly robust construction for generative training objectives, with instantiations including diffusion models and flow matching. Dance et al. [8 ] learn a flow from the observational treatment-conditional distribution to the interventional distribution, thereby exploiting geometric structure shared by the observed and target laws rather than generating the latter from a fixed generic base distribution. A common limitation of these approaches is that sampling from each interventional distribution of interest requires fitting a corresponding generative model, in addition to estimating the causal nuisances. This can be especially costly in high-dimensional settings, when multiple interventional distributions are of interest. IMMD-Flow instead avoids target-model training by formulating sampling as distributional optimization. For MMD, the target distribution is represented by a single kernel mean embedding, which admits a closed-form orthogonal estimator from observational data. Sampling then reduces to directly optimizing a particle approximation toward this estimated embedding.
Kernel embeddings of interventional distributions. Our target representation builds directly on work representing causal outcome distributions by kernel mean embeddings. Muandet et al. [12 ] introduced counterfactual mean embeddings (CMEs), which identify interventional outcome distributions nonparametrically in an RKHS and permit distributional comparison as well as sample reconstruction by kernel herding. Subsequent work developed doubly robust versions of these embeddings for distributional treatment-effect testing [43 ] , while Luedtke and Chung [29 ] developed general one-step theory for Hilbert-valued causal parameters, including counterfactual kernel mean embeddings. More recently, Zenati et al. [13 ] introduced counterfactual policy mean embeddings for distributional treatment effects under soft interventions, deriving doubly robust estimators for general target policies together with tests (see also [44 ] ) and a kernel-herding procedure for generating representative counterfactual samples [28 ] . Our method uses the same doubly robust kernel embedding, but as a fixed target for distributional optimization. Instead of kernel herding, IMMD-Flow jointly transports particles under the MMD gradient flow, allowing initialization from a useful source distribution such as P Y ∣ A = a \mathbb{P}_{Y\mid A=a} P Y ∣ A = a
MMD-based particle transport. Methodologically, our sampling stage is related to nonparametric generative methods based on kernel discrepancies and Wasserstein gradient flows. MMD provides a metric on probability measures for characteristic kernels [9 , 10 ] , and Arbel et al. [11 ] introduced the Wasserstein gradient flow of MMD together with interacting-particle approximations and noise regularization. Subsequent work has studied alternative kernels, regularization schemes, and scalable or deep versions of MMD transport [25 , 26 , 27 ] . Recent analyses have further clarified convergence properties of MMD and energy-kernel gradient flows. These methods assume that the target distribution, or samples from it, are available. Our setting introduces an additional causal-estimation layer: samples from the target interventional law are unobserved, and the force driving the particle dynamics is instead constructed from an estimated interventional mean embedding. The main connection between the two parts of our method is therefore that causal estimation error appears as a perturbation of the target potential in the MMD flow. This permits us to combine orthogonal estimation of the interventional distribution with particle-level transport guarantees in a single error analysis.
Counterfactual sampling. Our work is also related to the broader literature on sampling from counterfactual distributions, including structural causal models and counterfactual transport methods [45 , 46 , 47 , 48 ] . A key distinction is that IMMD-Flow targets the marginal interventional law P π \mathbb{P}^{\pi} P π unit-level counterfactual distributions such as P π ( ⋅ ∣ Y = y ) \mathbb{P}^{\pi}(\cdot\mid Y=y) P π ( ⋅ ∣ Y = y ) P Y ∣ A = a \mathbb{P}_{Y\mid A=a} P Y ∣ A = a P π P^\pi P π P π P^\pi P π
Orthogonal Target Computation
This appendix derives the orthogonal Hilbert-space target used by IMMD-Flow. We first give the full kernel-ridge computation, then derive its input-projected Nyström approximation. We subsequently describe the Monte Carlo approximation of the intervention integrals and compare the computational costs of the two constructions. Throughout, the target-policy integral is treated separately from its finite-J J J
Full-ridge computation
Write S i = ( A i , X i ) S_i=(A_i,X_i) S i = ( A i , X i ) A × X \mathcal{A}\times\mathcal{X} A × X
k S ( ( a , x ) , ( a ′ , x ′ ) ) = k A ( a , a ′ ) k X ( x , x ′ ) . k_S\left((a,x),(a',x')\right)=k_A(a,a')k_X(x,x'). k S ( ( a , x ) , ( a ′ , x ′ ) ) = k A ( a , a ′ ) k X ( x , x ′ ) . Let K ∈ R n × n K\in\mathbb{R}^{n\times n} K ∈ R n × n K i j = k S ( S i , S j ) K_{ij}=k_S(S_i,S_j) K ij = k S ( S i , S j )
B λ : = ( K + n λ I n ) − 1 . B_\lambda:= (K+n\lambda I_n)^{-1}. B λ := ( K + nλ I n ) − 1 . For s = ( a , x ) s=(a,x) s = ( a , x )
k s = ( k S ( S 1 , s ) , … , k S ( S n , s ) ) T . k_s=\left(k_S(S_1,s),\ldots,k_S(S_n,s)\right)^\mathsf{T}. k s = ( k S ( S 1 , s ) , … , k S ( S n , s ) ) T . The vector-valued kernel-ridge estimator of the conditional mean embedding is
m ^ f u l l ( s ) = ∑ j = 1 n [ B λ k s ] j φ ( Z j ) . (24) \widehat{m}_{\mathrm{full}}(s)=\sum_{j=1}^{n}[B_\lambda k_s]_j\varphi(Z_j). \tag*{(24)} m full ( s ) = j = 1 ∑ n [ B λ k s ] j φ ( Z j ) . (24) For each observed context X i X_i X i
k ‾ π , i : = ∫ A k ( a , X i ) π ( d a ∣ X i ) ∈ R n , (25) \overline{k}_{\pi,i}:=\int_{\mathcal{A}} k(a,X_i)\,\pi(da\mid X_i)\in\mathbb{R}^{n}, \tag*{(25)} k π , i := ∫ A k ( a , X i ) π ( d a ∣ X i ) ∈ R n , (25) and collect these vectors as
K π : = ( k ‾ π , 1 , … , k ‾ π , n ) ∈ R n × n . K_\pi:=\left(\overline{k}_{\pi,1},\ldots,\overline{k}_{\pi,n}\right)\in\mathbb{R}^{n\times n}. K π := ( k π , 1 , … , k π , n ) ∈ R n × n . For the product kernel above, this matrix can equivalently be written as
K π = K X ⊙ K A , π , K_\pi=K_X\odot K_{A,\pi}, K π = K X ⊙ K A , π , where
[ K A , π ] j i : = ∫ A k A ( A j , a ) π ( d a ∣ X i ) . (26) [K_{A,\pi}]_{ji}:=\int_{\mathcal{A}} k_A(A_j,a)\,\pi(da\mid X_i). \tag*{(26)} [ K A , π ] j i := ∫ A k A ( A j , a ) π ( d a ∣ X i ) . (26) Writing
w ^ = ( w ^ ( A 1 , X 1 ) , … , w ^ ( A n , X n ) ) T , \widehat{w}=\left(\widehat{w}(A_1,X_1),\ldots,\widehat{w}(A_n,X_n)\right)^\mathsf{T}, w = ( w ( A 1 , X 1 ) , … , w ( A n , X n ) ) T , substitution of equation 20 into the empirical orthogonal score gives
h ^ π , f u l l = ∑ i = 1 n v f u l l , i φ ( Z i ) , (27) \widehat{h}_{\pi,\mathrm{full}}=\sum_{i=1}^{n}v_{\mathrm{full},i}\varphi(Z_i), \tag*{(27)} h π , full = i = 1 ∑ n v full , i φ ( Z i ) , (27) where
b f u l l = B λ ( K π 1 n − K w ^ ) , v f u l l = 1 n ( b f u l l + w ^ ) . (28) b^{\mathrm{full}}=B_\lambda\left(K_\pi\mathbf{1}_n-K\widehat{w}\right),\quad v_{\mathrm{full}}=\frac{1}{n}\left(b^{\mathrm{full}}+\widehat{w}\right). \tag*{(28)} b full = B λ ( K π 1 n − K w ) , v full = n 1 ( b full + w ) . (28) Equivalently, defining the full-ridge plug-in coefficients
v P I , f u l l : = B λ K π 1 n n , (29) v_{\mathrm{PI},\mathrm{full}}:=B_\lambda K_\pi\frac{\mathbf{1}_n}{n}, \tag*{(29)} v PI , full := B λ K π n 1 n , (29) the doubly robust coefficients satisfy
v f u l l = v P I , f u l l + 1 n ( w ^ − B λ K w ^ ) . (30) v_{\mathrm{full}} = v_{\mathrm{PI,full}} + \frac{1}{n}\left(\widehat{w} - B_{\lambda}K\widehat{w}\right). \tag*{(30)} v full = v PI , full + n 1 ( w − B λ K w ) . (30) This is the finite-sample full-kernel form of the doubly robust policy mean embedding of Zenati et al. [13 ] . Thus, once the regression system is solved, changing the intervention affects the estimator only through the intervention integrals and, when target-specific, the density-ratio nuisance.
Projected Nyström computation
The full construction requires solving an n × n n \times n n × n [31 , 49 ] .
Select landmark indices
L = { ℓ 1 , … , ℓ m L } ⊆ { 1 , … , n } L = \{\ell_1,\ldots,\ell_{m_L}\} \subseteq\{1,\ldots,n\} L = { ℓ 1 , … , ℓ m L } ⊆ { 1 , … , n } uniformly without replacement, and write
S ~ q : = S ℓ q . \widetilde{S}_q := S_{\ell_q}. S q := S ℓ q . Define
C i q : = k S ( S i , S ~ q ) , W q r : = k S ( S ~ q , S ~ r ) , (31) C_{iq} := k_S(S_i,\widetilde{S}_q), \qquad W_{qr} := k_S(\widetilde{S}_q,\widetilde{S}_r), \tag*{(31)} C i q := k S ( S i , S q ) , W q r := k S ( S q , S r ) , (31) so that
C ∈ R n × m L , W ∈ R m L × m L . C \in\mathbb{R}^{n \times m_L}, \qquad W \in\mathbb{R}^{m_L \times m_L}. C ∈ R n × m L , W ∈ R m L × m L . Let
G λ : = ( C T C + n λ W ) † , (32) G_{\lambda} := \left(C^{\mathsf{T}}C + n\lambda W\right)^{\dagger}, \tag*{(32)} G λ := ( C T C + nλW ) † , (32) where † \dagger †
For s = ( a , x ) s=(a,x) s = ( a , x )
c s = ( k S ( S ~ 1 , s ) , … , k S ( S ~ m L , s ) ) T . c_s = \left(k_S(\widetilde{S}_1,s),\ldots,k_S(\widetilde{S}_{m_L},s)\right)^{\mathsf{T}}. c s = ( k S ( S 1 , s ) , … , k S ( S m L , s ) ) T . The projected estimator can then be written directly in the logged-outcome basis as
m ^ N y s ( s ) = ∑ i = 1 n [ C G λ c s ] i ϕ ( Z i ) . (33) \widehat{m}_{\mathrm{Nys}}(s) = \sum_{i=1}^{n}[CG_{\lambda}c_s]_i\phi(Z_i). \tag*{(33)} m Nys ( s ) = i = 1 ∑ n [ C G λ c s ] i ϕ ( Z i ) . (33) For each context X i X_i X i
c ‾ π , i : = ∫ A c ( a , X i ) π ( d a ∣ X i ) ∈ R m L , (34) \overline{c}_{\pi,i} := \int_{\mathcal{A}} c(a,X_i)\,\pi(\mathrm{d}a \mid X_i) \in\mathbb{R}^{m_L}, \tag*{(34)} c π , i := ∫ A c ( a , X i ) π ( d a ∣ X i ) ∈ R m L , (34) and let
C π : = ( c ‾ π , 1 , … , c ‾ π , n ) ∈ R m L × n . C_{\pi} := \left(\overline{c}_{\pi,1},\ldots,\overline{c}_{\pi,n}\right) \in\mathbb{R}^{m_L \times n}. C π := ( c π , 1 , … , c π , n ) ∈ R m L × n . Because the landmarks are selected from the logged inputs,
C π = ( K π ) L : . C_{\pi} = (K_{\pi})_{L:}. C π = ( K π ) L : . Thus the projected construction evaluates the exact intervention-integrated regression contrast only at the landmarks.
Substituting equation (33) into the orthogonal score gives
b p r o j = C G λ ( C π 1 n − C T w ^ ) , v D R = 1 n ( b p r o j + w ^ ) . (35) b^{\mathrm{proj}} = CG_{\lambda}\left(C_{\pi}\mathbf{1}_n - C^{\mathsf{T}}\widehat{w}\right), \qquad v_{\mathrm{DR}} = \frac{1}{n}\left(b^{\mathrm{proj}}+\widehat{w}\right). \tag*{(35)} b proj = C G λ ( C π 1 n − C T w ) , v DR = n 1 ( b proj + w ) . (35) Hence the exact-policy-integral projected target is
h ^ π , N y s = ∑ i = 1 n v D R , i ϕ ( Z i ) . (36) \widehat{h}_{\pi,\mathrm{Nys}} = \sum_{i=1}^{n}v_{\mathrm{DR},i}\phi(Z_i). \tag*{(36)} h π , Nys = i = 1 ∑ n v DR , i ϕ ( Z i ) . (36) The projection acts on the input-side regression problem only: equation (36) still expands over all n n n
Whitened landmark coordinates. For numerical stability, let
W = V r D r V r ⊤ W = V_r D_r V_r^{\top} W = V r D r V r ⊤ denote the eigendecomposition of W W W
R W : = V r D r − 1 / 2 , Φ : = C R W . R_W := V_r D_r^{-1/2}, \qquad\Phi:= C R_W. R W := V r D r − 1/2 , Φ := C R W . Writing
r π : = C π 1 n − C ⊤ w ^ , q π : = R W ⊤ r π , (37) r_\pi:= C_\pi\mathbf{1}_n - C^{\top}\widehat{w}, \qquad q_\pi:= R_W^{\top}r_\pi, \tag*{(37)} r π := C π 1 n − C ⊤ w , q π := R W ⊤ r π , (37) the projected coefficients can equivalently be computed as
b p r o j = Φ ( Φ ⊤ Φ + n λ I r ) − 1 q π . (38) b^{\mathrm{proj}} = \Phi\left(\Phi^{\top}\Phi+n\lambda I_r\right)^{-1}q_\pi. \tag*{(38)} b proj = Φ ( Φ ⊤ Φ + nλ I r ) − 1 q π . (38) Using all positive eigenpairs is only a change of coordinates within the landmark subspace. Discarding additional positive eigenvalues through a numerical threshold introduces a further subspace approximation.
Intervention integration and computational reuse
When the intervention integrals in equation (25) and equation (34) are unavailable analytically, we draw
A i 1 π , … , A i J π ∼ i i d π ( ⋅ ∣ X i ) A_{i1}^{\pi},\ldots,A_{iJ}^{\pi}\overset{\mathrm{iid}}{\sim}\pi(\cdot\mid X_i) A i 1 π , … , A i J π ∼ iid π ( ⋅ ∣ X i ) for every logged context X i X_i X i
For the full-ridge construction, define
k ~ π , i ( J ) : = 1 J ∑ j = 1 J k ( A i j π , X i ) ∈ R n , (39) \widetilde{k}_{\pi,i}^{(J)} := \frac{1}{J}\sum_{j=1}^{J} k\left(A_{ij}^{\pi},X_i\right)\in\mathbb{R}^{n}, \tag*{(39)} k π , i ( J ) := J 1 j = 1 ∑ J k ( A ij π , X i ) ∈ R n , (39) and collect these vectors as K ~ π ( J ) : = ( k ~ π , 1 ( J ) , … , k ~ π , n ( J ) ) \widetilde{K}_{\pi}^{(J)} := \left(\widetilde{k}_{\pi,1}^{(J)},\ldots,\widetilde{k}_{\pi,n}^{(J)}\right) K π ( J ) := ( k π , 1 ( J ) , … , k π , n ( J ) ) J J J
b ~ J f u l l : = B λ ( K ~ π ( J ) 1 n − K w ^ ) , v ~ f u l l , J : = 1 n ( b ~ J f u l l + w ^ ) , (40) \widetilde{b}_{J}^{\mathrm{full}} := B_\lambda\left(\widetilde{K}_{\pi}^{(J)}\mathbf{1}_n-K\widehat{w}\right), \qquad\widetilde{v}_{\mathrm{full},J} := \frac{1}{n}\left(\widetilde{b}_{J}^{\mathrm{full}}+\widehat{w}\right), \tag*{(40)} b J full := B λ ( K π ( J ) 1 n − K w ) , v full , J := n 1 ( b J full + w ) , (40) with numerical target
h ~ π , f u l l , J : = ∑ i = 1 n v ~ f u l l , J , i φ ( Z i ) . (41) \widetilde{h}_{\pi,\mathrm{full},J} := \sum_{i=1}^{n}\widetilde{v}_{\mathrm{full},J,i}\varphi(Z_i). \tag*{(41)} h π , full , J := i = 1 ∑ n v full , J , i φ ( Z i ) . (41) For the projected construction, use
c ~ π , i ( J ) : = 1 J ∑ j = 1 J c ( A i j π , X i ) . (42) \widetilde{c}_{\pi,i}^{(J)} := \frac{1}{J}\sum_{j=1}^{J} c\left(A_{ij}^{\pi},X_i\right). \tag*{(42)} c π , i ( J ) := J 1 j = 1 ∑ J c ( A ij π , X i ) . (42) Let
C ~ π ( J ) : = ( c ~ π , 1 ( J ) , … , c ~ π , n ( J ) ) . \widetilde{C}_{\pi}^{(J)} := \left(\widetilde{c}_{\pi,1}^{(J)},\ldots,\widetilde{c}_{\pi,n}^{(J)}\right). C π ( J ) := ( c π , 1 ( J ) , … , c π , n ( J ) ) . Using the same target-policy draws in the full and projected constructions,
C ~ π ( J ) = ( K ~ π ( J ) ) L : . \widetilde{C}_{\pi}^{(J)}=\left(\widetilde{K}_{\pi}^{(J)}\right)_{L:}. C π ( J ) = ( K π ( J ) ) L : . Replacing C π C_\pi C π C ~ π ( J ) \widetilde{C}_{\pi}^{(J)} C π ( J ) (35) gives
b ~ J p r o j : = C G λ ( C ~ π ( J ) 1 n − C ⊤ w ^ ) , v ~ D R , J : = 1 n ( b ~ J p r o j + w ^ ) . (43) \widetilde{b}_{J}^{\mathrm{proj}} := CG_\lambda\left(\widetilde{C}_{\pi}^{(J)}\mathbf{1}_n-C^{\top}\widehat{w}\right), \qquad\widetilde{v}_{\mathrm{DR},J} := \frac{1}{n}\left(\widetilde{b}_{J}^{\mathrm{proj}}+\widehat{w}\right). \tag*{(43)} b J proj := C G λ ( C π ( J ) 1 n − C ⊤ w ) , v DR , J := n 1 ( b J proj + w ) . (43) The resulting numerical target is
h ~ π , N y s , J : = ∑ i = 1 n v ~ D R , J , i φ ( Z i ) . (44) \widetilde{h}_{\pi,\mathrm{Nys},J} := \sum_{i=1}^{n}\widetilde{v}_{\mathrm{DR},J,i}\varphi(Z_i). \tag*{(44)} h π , Nys , J := i = 1 ∑ n v DR , J , i φ ( Z i ) . (44) This is the finite-J J J h ~ π \widetilde{h}_\pi h π
h ~ π = h ~ π , N y s , J . \widetilde{h}_\pi=\widetilde{h}_{\pi,\mathrm{Nys},J}. h π = h π , Nys , J . Similarly, h ^ π = h ^ π , N y s \widehat{h}_{\pi}=\widehat{h}_{\pi,\mathrm{Nys}} h π = h π , Nys
Let F n \mathcal{F}_n F n F n \mathcal{F}_n F n
E [ K ~ π ( J ) ∣ F n ] = K π , E [ C ~ π ( J ) ∣ F n ] = C π . (45) \mathbb{E}\left[\widetilde{K}_{\pi}^{(J)} \mid\mathcal{F}_n\right] = K_{\pi}, \qquad
\mathbb{E}\left[\widetilde{C}_{\pi}^{(J)} \mid\mathcal{F}_n\right] = C_{\pi}. \tag*{(45)} E [ K π ( J ) ∣ F n ] = K π , E [ C π ( J ) ∣ F n ] = C π . (45) Because both coefficient maps are linear in the intervention-integrated right-hand side,
E [ h ^ π , f u l l , J ∣ F n ] = h ^ π , f u l l , E [ h ^ π , N y s , J ∣ F n ] = h ^ π , N y s . (46) \mathbb{E}\left[\widehat{h}_{\pi,\mathrm{full},J} \mid\mathcal{F}_n\right] = \widehat{h}_{\pi,\mathrm{full}}, \qquad
\mathbb{E}\left[\widehat{h}_{\pi,\mathrm{Nys},J} \mid\mathcal{F}_n\right] = \widehat{h}_{\pi,\mathrm{Nys}}. \tag*{(46)} E [ h π , full , J ∣ F n ] = h π , full , E [ h π , Nys , J ∣ F n ] = h π , Nys . (46) The corresponding Hilbert-space Monte Carlo error is analyzed in Appendix § 14 .
For fixed representation, kernels, and adjustment variables, the landmark set, matrices C C C W W W C T C + n λ W C^{\mathsf{T}}C+n\lambda W C T C + nλW
Thus, evaluating a new supported intervention does not require refitting the conditional mean embedding. It requires only the intervention-specific policy integrals, density-ratio evaluations, and the resulting linear-system right-hand side.
Plug-in and inverse-probability variants
Our benchmark tables also compare two single-nuisance IMMD variants. For the finite-J J J
v ~ P I , J : = 1 n C G λ C ~ π ( J ) 1 n , v ~ I P W : = 1 n w ^ , \widetilde{v}_{\mathrm{PI},J} := \frac{1}{n}CG_{\lambda}\widetilde{C}_{\pi}^{(J)}\mathbf{1}_n,\qquad
\widetilde{v}_{\mathrm{IPW}} := \frac{1}{n}\widehat{w}, v PI , J := n 1 C G λ C π ( J ) 1 n , v IPW := n 1 w , v ~ D R , J : = v ~ P I , J + 1 n ( w ^ − C G λ C T w ^ ) . \widetilde{v}_{\mathrm{DR},J} := \widetilde{v}_{\mathrm{PI},J}+\frac{1}{n}\left(\widehat{w}-CG_{\lambda}C^{\mathsf{T}}\widehat{w}\right). v DR , J := v PI , J + n 1 ( w − C G λ C T w ) . Thus, Plug-in–IMMD uses only the fitted conditional mean embedding, IPW–IMMD uses only the estimated policy ratio, and DR–IMMD combines both. The displayed IPW vector is the Horvitz–Thompson form; any declared benchmark-specific unit-mass rescaling is applied after these raw coefficients are constructed. For e ∈ { P I , I P W , D R } e \in\{\mathrm{PI},\mathrm{IPW},\mathrm{DR}\} e ∈ { PI , IPW , DR } v ~ e \widetilde{v}_e v e J J J
h ^ π , e = ∑ i = 1 n v ~ e , i φ ( Z i ) \widehat{h}_{\pi,e}=\sum_{i=1}^n\widetilde{v}_{e,i}\varphi(Z_i) h π , e = ∑ i = 1 n v e , i φ ( Z i )
and apply the same particle flow. Hence the variants differ only in the target coefficients supplied to the optimizer.
Computational complexity
We compare the time and memory costs of the full-ridge and input-projected Nyström target constructions. Table 3 summarizes the comparison, while the discussion below details the dominant operations and the computations that can be reused across interventions.
Table 3. Computational comparison of the full-ridge and input-projected Nyström target constructions. The per-intervention costs assume that the corresponding regression system has already been factorized.
For the full-ridge estimator, constructing and factorizing the n × n n \times n n × n O ( n 3 ) O(n^3) O ( n 3 ) O ( n 2 ) O(n^2) O ( n 2 ) J J J O ( J n 2 ) O(Jn^2) O ( J n 2 ) O ( n 2 ) O(n^2) O ( n 2 ) C ⊤ C C^{\top}C C ⊤ C O ( n m L 2 ) O(nm_L^2) O ( n m L 2 ) m L × m L m_L \times m_L m L × m L O ( m L 3 ) O(m_L^3) O ( m L 3 ) C C C W W W O ( n m L + m L 2 ) O(nm_L + m_L^2) O ( n m L + m L 2 ) J J J O ( J m L n ) O(Jm_Ln) O ( J m L n ) O ( n m L + m L 2 ) O(nm_L + m_L^2) O ( n m L + m L 2 ) J J J O ( n m L 2 + m L 3 ) O(nm_L^2 + m_L^3) O ( n m L 2 + m L 3 ) O ( n m L + m L 2 ) O(nm_L + m_L^2) O ( n m L + m L 2 )
Input-side Nyström approximation reduces the cost of conditional-embedding regression and target construction, but it does not reduce the number of logged outcome features in the final target expansion. Consequently, direct evaluation of the target contribution for M M M O ( M n ) O(Mn) O ( M n )
If the intervention integrals are available analytically, the factors proportional to J J J
The statistical guarantees in Appendix § 14 require the landmark count m L m_L m L n n n
Proofs for the Interventional MMD Construction
Double robustness and orthogonality
Proof of Proposition 3.2. By the change-of-measure identity,
E 0 [ ∫ A m ( a , X ) π ( d a ∣ X ) ] = E 0 [ w 0 ( A , X ) m ( A , X ) ] . (47) \mathbb{E}_0\left[\int_{\mathcal{A}} m(a,X)\,\pi(da\mid X)\right] = \mathbb{E}_0[w_0(A,X)m(A,X)]. \tag*{(47)} E 0 [ ∫ A m ( a , X ) π ( d a ∣ X ) ] = E 0 [ w 0 ( A , X ) m ( A , X )] . (47) Moreover, E 0 [ φ ( Z ) ∣ A , X ] = m 0 ( A , X ) \mathbb{E}_0[\varphi(Z)\mid A,X] = m_0(A,X) E 0 [ φ ( Z ) ∣ A , X ] = m 0 ( A , X )
h η = E 0 [ w 0 m + w ( m 0 − m ) ] . h_\eta= \mathbb{E}_0[w_0m + w(m_0-m)]. h η = E 0 [ w 0 m + w ( m 0 − m )] . Since χ ( π ) = E 0 [ w 0 m 0 ] \chi(\pi) = \mathbb{E}_0[w_0m_0] χ ( π ) = E 0 [ w 0 m 0 ]
h η − χ ( π ) = E 0 [ w 0 m + w ( m 0 − m ) − w 0 m 0 ] = − E 0 [ ( w − w 0 ) ( m − m 0 ) ] . h_\eta-\chi(\pi) = \mathbb{E}_0[w_0m + w(m_0-m)-w_0m_0]
= -\mathbb{E}_0[(w-w_0)(m-m_0)]. h η − χ ( π ) = E 0 [ w 0 m + w ( m 0 − m ) − w 0 m 0 ] = − E 0 [( w − w 0 ) ( m − m 0 )] . Taking norms and applying the Hilbert-valued Cauchy–Schwarz inequality gives
∥ h η − χ ( π ) ∥ H ≤ ∥ w − w 0 ∥ L 2 ( P A , X ) ∥ m − m 0 ∥ L 2 ( P A , X ; H ) = r w r m . \lVert h_\eta-\chi(\pi)\rVert_{\mathcal{H}}
\le\lVert w-w_0\rVert_{L_2(P_{A,X})}\lVert m-m_0\rVert_{L_2(P_{A,X};\mathcal{H})}
= r_w r_m. ∥ h η − χ ( π ) ∥ H ≤ ∥ w − w 0 ∥ L 2 ( P A , X ) ∥ m − m 0 ∥ L 2 ( P A , X ; H ) = r w r m . The double-robustness statements follow immediately.
For a direction η ˙ = ( m ˙ , w ˙ ) \dot{\eta}=(\dot{m},\dot{w}) η ˙ = ( m ˙ , w ˙ )
η ϵ = η 0 + ϵ η ˙ = ( m 0 + ϵ m ˙ , w 0 + ϵ w ˙ ) . \eta_\epsilon= \eta_0+\epsilon\dot{\eta}=(m_0+\epsilon\dot{m},w_0+\epsilon\dot{w}). η ϵ = η 0 + ϵ η ˙ = ( m 0 + ϵ m ˙ , w 0 + ϵ w ˙ ) . The remainder identity gives
h η ϵ − χ ( π ) = − ϵ 2 E 0 [ w ˙ m ˙ ] . h_{\eta_\epsilon}-\chi(\pi)=-\epsilon^2\mathbb{E}_0[\dot{w}\dot{m}]. h η ϵ − χ ( π ) = − ϵ 2 E 0 [ w ˙ m ˙ ] . Hence
D η h η 0 [ η ˙ ] = 0. D_\eta h_{\eta_0}[\dot{\eta}]=0. D η h η 0 [ η ˙ ] = 0. Since u u u η \eta η
D η ( u − h η ) ∣ η = η 0 [ η ˙ ] = − D η h η 0 [ η ˙ ] = 0 , \left.D_\eta(u-h_\eta)\right|_{\eta=\eta_0}[\dot{\eta}]
=-D_\eta h_{\eta_0}[\dot{\eta}]=0, D η ( u − h η ) ∣ η = η 0 [ η ˙ ] = − D η h η 0 [ η ˙ ] = 0 , which proves equation ORTH.
Dissipation and velocity stability
Proof of Proposition 3.5 . We first record the spatial gradient of the witness. For fixed Q ∈ P 2 ( Z ) Q\in\mathcal{P}_2(\mathcal{Z}) Q ∈ P 2 ( Z ) h ∈ H h\in\mathcal{H} h ∈ H u ∈ R d u\in\mathbb{R}^d u ∈ R d φ \varphi φ
D F Q , h ( z ) [ u ] = ⟨ μ Q − h , D φ ( z ) [ u ] ⟩ H = ⟨ D φ ( z ) ∗ ( μ Q − h ) , u ⟩ 2 . D_{F_{Q,h}(z)}[u]=\langle\mu_Q-h,D\varphi(z)[u]\rangle_{\mathcal{H}}=\langle D\varphi(z)^*(\mu_Q-h),u\rangle_2. D F Q , h ( z ) [ u ] = ⟨ μ Q − h , D φ ( z ) [ u ] ⟩ H = ⟨ D φ ( z ) ∗ ( μ Q − h ) , u ⟩ 2 . Consequently,
∇ f Q , h ( z ) = D φ ( z ) ∗ ( μ Q − h ) , v h ( Q , z ) = − D φ ( z ) ∗ ( μ Q − h ) . (48) \nabla f_{Q,h}(z)=D\varphi(z)^*(\mu_Q-h),\qquad v_h(Q,z)=-D\varphi(z)^*(\mu_Q-h). \tag*{(48)} ∇ f Q , h ( z ) = D φ ( z ) ∗ ( μ Q − h ) , v h ( Q , z ) = − D φ ( z ) ∗ ( μ Q − h ) . (48) Let
μ t : = μ Q t h = ∫ Z φ ( z ) Q t h ( d z ) , v t ( z ) : = v h ( Q t h , z ) . \mu_t:=\mu_{Q_t^h}=\int_{\mathcal{Z}}\varphi(z)Q_t^h(dz),\qquad v_t(z):=v_h(Q_t^h,z). μ t := μ Q t h = ∫ Z φ ( z ) Q t h ( d z ) , v t ( z ) := v h ( Q t h , z ) . To identify the derivative of μ t \mu_t μ t g ∈ H g\in\mathcal{H} g ∈ H
ψ g ( z ) : = ⟨ g , φ ( z ) ⟩ H . \psi_g(z):=\langle g,\varphi(z)\rangle_{\mathcal{H}}. ψ g ( z ) := ⟨ g , φ ( z ) ⟩ H . Its spatial gradient is
∇ ψ g ( z ) = D φ ( z ) ∗ g . \nabla\psi_g(z)=D\varphi(z)^*g. ∇ ψ g ( z ) = D φ ( z ) ∗ g . Applying the weak formulation of the continuity equation to ψ g \psi_g ψ g t t t
d d t ⟨ g , μ t ⟩ H = d d t ∫ Z ψ g ( z ) Q t h ( d z ) = ∫ Z ⟨ D φ ( z ) ∗ g , v t ( z ) ⟩ 2 Q t h ( d z ) = ⟨ g , ∫ Z D φ ( z ) [ v t ( z ) ] Q t h ( d z ) ⟩ H . \begin{aligned}
\frac{d}{dt}\langle g,\mu_t\rangle_{\mathcal{H}}
&=\frac{d}{dt}\int_{\mathcal{Z}}\psi_g(z)Q_t^h(dz)\\
&=\int_{\mathcal{Z}}\langle D\varphi(z)^*g,v_t(z)\rangle_2Q_t^h(dz)\\
&=\left\langle g,\int_{\mathcal{Z}}D\varphi(z)[v_t(z)]Q_t^h(dz)\right\rangle_{\mathcal{H}}.
\end{aligned} d t d ⟨ g , μ t ⟩ H = d t d ∫ Z ψ g ( z ) Q t h ( d z ) = ∫ Z ⟨ D φ ( z ) ∗ g , v t ( z ) ⟩ 2 Q t h ( d z ) = ⟨ g , ∫ Z D φ ( z ) [ v t ( z )] Q t h ( d z ) ⟩ H . Since this holds for every g ∈ H g\in\mathcal{H} g ∈ H
μ ˙ t = ∫ Z D φ ( z ) [ v t ( z ) ] Q t h ( d z ) (49) \dot{\mu}_t=\int_{\mathcal{Z}}D\varphi(z)[v_t(z)]Q_t^h(dz) \tag*{(49)} μ ˙ t = ∫ Z D φ ( z ) [ v t ( z )] Q t h ( d z ) (49) for almost every t t t D φ D\varphi D φ v t ∈ L 2 ( Q t h ; R d ) v_t\in L_2(Q_t^h;\mathbb{R}^d) v t ∈ L 2 ( Q t h ; R d )
Using the Hilbert-space chain rule and equation (49) ,
d d t F h ( Q t h ) = d d t 1 2 ∥ μ t − h ∥ H 2 = ⟨ μ t − h , μ ˙ t ⟩ H = ∫ Z ⟨ D φ ( z ) ∗ ( μ t − h ) , v t ( z ) ⟩ 2 Q t h ( d z ) . \begin{aligned}
\frac{d}{dt}F_h(Q_t^h)
&=\frac{d}{dt}\frac{1}{2}\|\mu_t-h\|_{\mathcal{H}}^2\\
&=\langle\mu_t-h,\dot{\mu}_t\rangle_{\mathcal{H}}\\
&=\int_{\mathcal{Z}}\langle D\varphi(z)^*(\mu_t-h),v_t(z)\rangle_2Q_t^h(dz).
\end{aligned} d t d F h ( Q t h ) = d t d 2 1 ∥ μ t − h ∥ H 2 = ⟨ μ t − h , μ ˙ t ⟩ H = ∫ Z ⟨ D φ ( z ) ∗ ( μ t − h ) , v t ( z ) ⟩ 2 Q t h ( d z ) . By equation (48) ,
D φ ( z ) ∗ ( μ t − h ) = − v t ( z ) . D\varphi(z)^*(\mu_t-h)=-v_t(z). D φ ( z ) ∗ ( μ t − h ) = − v t ( z ) . Therefore,
d d t F h ( Q t h ) = − ∫ Z ∥ v h ( Q t h , z ) ∥ 2 2 Q t h ( d z ) ≤ 0 , (50) \frac{d}{dt}F_h(Q_t^h)=-\int_{\mathcal{Z}}\|v_h(Q_t^h,z)\|_2^2Q_t^h(dz)\leq0, \tag*{(50)} d t d F h ( Q t h ) = − ∫ Z ∥ v h ( Q t h , z ) ∥ 2 2 Q t h ( d z ) ≤ 0 , (50) which proves the dissipation identity.
For velocity stability, let h 1 , h 2 ∈ H h_1,h_2\in\mathcal{H} h 1 , h 2 ∈ H (48) ,
v h 1 ( Q , z ) − v h 2 ( Q , z ) = D φ ( z ) ∗ ( h 1 − h 2 ) . v_{h_1}(Q,z)-v_{h_2}(Q,z)=D\varphi(z)^*(h_1-h_2). v h 1 ( Q , z ) − v h 2 ( Q , z ) = D φ ( z ) ∗ ( h 1 − h 2 ) . Hence
∥ v h 1 ( Q , z ) − v h 2 ( Q , z ) ∥ 2 ≤ ∥ D φ ( z ) ∥ o p ∥ h 1 − h 2 ∥ H ≤ L φ ∥ h 1 − h 2 ∥ H . \|v_{h_1}(Q,z)-v_{h_2}(Q,z)\|_2\leq\|D\varphi(z)\|_{\mathrm{op}}\|h_1-h_2\|_{\mathcal{H}}\leq L_\varphi\|h_1-h_2\|_{\mathcal{H}}. ∥ v h 1 ( Q , z ) − v h 2 ( Q , z ) ∥ 2 ≤ ∥ D φ ( z ) ∥ op ∥ h 1 − h 2 ∥ H ≤ L φ ∥ h 1 − h 2 ∥ H . Taking the supremum over z ∈ Z z\in\mathcal{Z} z ∈ Z
h 1 = h η , h 2 = χ ( π ) , h_1=h_\eta,\qquad h_2=\chi(\pi), h 1 = h η , h 2 = χ ( π ) , gives
sup z ∈ Z ∥ v h η ( Q , z ) − v χ ( π ) ( Q , z ) ∥ 2 ≤ L φ ∥ h η − χ ( π ) ∥ H ≤ L φ r w r m , \sup_{z\in\mathcal{Z}}\|v_{h_\eta}(Q,z)-v_{\chi(\pi)}(Q,z)\|_2\leq L_\varphi\|h_\eta-\chi(\pi)\|_{\mathcal{H}}\leq L_\varphi r_w r_m, z ∈ Z sup ∥ v h η ( Q , z ) − v χ ( π ) ( Q , z ) ∥ 2 ≤ L φ ∥ h η − χ ( π ) ∥ H ≤ L φ r w r m , where the final inequality follows from Proposition 3.2.
Theoretical Analysis
Projected Nyström approximation of the conditional mean embedding
This section controls the error introduced when the conditional mean embedding nuisance is estimated by input-projected Nyström kernel ridge regression. The argument specializes the vector-valued regression analysis of El Ahmad et al. [49 ] to our setting. In contrast to their general construction, we sketch only the input RKHS. The latent-output RKHS remains unchanged, so no output-projection approximation term appears.
Statistical setup. Let S : = A × X \mathcal{S} := \mathcal{A} \times\mathcal{X} S := A × X S : = ( A , X ) S := (A, X) S := ( A , X ) Z : = E ( Y ) Z := E(Y) Z := E ( Y ) E E E ( S i , Z i ) i = 1 n (S_i, Z_i)_{i=1}^{n} ( S i , Z i ) i = 1 n P 0 P_0 P 0
Let k S : S × S → R k_S : \mathcal{S} \times\mathcal{S} \to\mathbb{R} k S : S × S → R H S \mathcal{H}_S H S
φ S ( s ) : = k S ( s , ⋅ ) . \varphi_S(s) := k_S(s, \cdot). φ S ( s ) := k S ( s , ⋅ ) . Let k Z : Z × Z → R k_Z : \mathcal{Z} \times\mathcal{Z} \to\mathbb{R} k Z : Z × Z → R H Z \mathcal{H}_Z H Z
φ Z ( z ) : = k Z ( z , ⋅ ) . \varphi_Z(z) := k_Z(z, \cdot). φ Z ( z ) := k Z ( z , ⋅ ) . Assumption D.1 (Measurability and kernel regularity). We assume that
The context, action, and latent spaces X \mathcal{X} X A \mathcal{A} A Z \mathcal{Z} Z σ \sigma σ
The RKHSs H S \mathcal{H}_S H S H Z \mathcal{H}_Z H Z
There exist finite constants κ S , κ Z > 0 \kappa_S, \kappa_Z > 0 κ S , κ Z > 0
sup s ∈ S k S ( s , s ) ≤ κ S 2 , sup z ∈ Z k Z ( z , z ) ≤ κ Z 2 . \sup_{s \in\mathcal{S}} k_S(s,s) \le\kappa_S^2,\qquad\sup_{z \in\mathcal{Z}} k_Z(z,z) \le\kappa_Z^2. s ∈ S sup k S ( s , s ) ≤ κ S 2 , z ∈ Z sup k Z ( z , z ) ≤ κ Z 2 . Define the latent conditional mean embedding
m 0 ( s ) : = E P 0 [ φ Z ( Z ) ∣ S = s ] ∈ H Z . m_0(s) := \mathbb{E}_{P_0}[\varphi_Z(Z) \mid S=s] \in\mathcal{H}_Z. m 0 ( s ) := E P 0 [ φ Z ( Z ) ∣ S = s ] ∈ H Z . We use the identity-decomposable operator-valued kernel (OVK)
K ( s , s ′ ) = k S ( s , s ′ ) I H Z . K(s,s') = k_S(s,s')I_{\mathcal{H}_Z}. K ( s , s ′ ) = k S ( s , s ′ ) I H Z . Its vector-valued RKHS is isometrically identified with the space S 2 ( H S , H Z ) \mathcal{S}_2(\mathcal{H}_S,\mathcal{H}_Z) S 2 ( H S , H Z ) H : H S → H Z H : \mathcal{H}_S \to\mathcal{H}_Z H : H S → H Z
m H ( s ) = H φ S ( s ) , ∥ m H ∥ G = ∥ H ∥ H S . m_H(s) = H\varphi_S(s),\qquad\|m_H\|_{\mathcal{G}} = \|H\|_{\mathrm{HS}}. m H ( s ) = H φ S ( s ) , ∥ m H ∥ G = ∥ H ∥ HS . Assumption D.2 (Attainability; El Ahmad et al. [49 Assumption 1] ). There exists H 0 ∈ S 2 ( H S , H Z ) H_0 \in\mathcal{S}_2(\mathcal{H}_S,\mathcal{H}_Z) H 0 ∈ S 2 ( H S , H Z ) such that
m 0 ( s ) = H 0 φ S ( s ) P S , 0 -almost surely . m_0(s) = H_0\varphi_S(s)\qquad P_{S,0}\text{-almost surely}. m 0 ( s ) = H 0 φ S ( s ) P S , 0 -almost surely . We now define the covariance operators associated with conditional mean embedding learning.
Definition D.3 (Population covariance operators). Let P S , 0 P_{S,0} P S , 0 S = ( A , X ) S = (A,X) S = ( A , X )
C S : = E P 0 [ φ S ( S ) ⊗ φ S ( S ) ] : H S → H S . C_S := \mathbb{E}_{P_0}[\varphi_S(S) \otimes\varphi_S(S)] : \mathcal{H}_S \to\mathcal{H}_S. C S := E P 0 [ φ S ( S ) ⊗ φ S ( S )] : H S → H S . and the population output–input cross-covariance operator
C Z S : = E P 0 [ φ Z ( Z ) ⊗ φ S ( S ) ] : H S → H Z . C_{ZS} := \mathbb{E}_{P_0}[\varphi_Z(Z) \otimes\varphi_S(S)] : \mathcal{H}_S \to\mathcal{H}_Z. C Z S := E P 0 [ φ Z ( Z ) ⊗ φ S ( S )] : H S → H Z . Under Assumption D.1, these Bochner expectations are well defined. The operator C S C_S C S
Tr ( C S ) = E ∥ φ S ( S ) ∥ H S 2 ≤ κ S 2 . \operatorname{Tr}(C_S)=\mathbb{E}\lVert\varphi_S(S)\rVert_{\mathcal{H}_S}^{2}\leq\kappa_S^2. Tr ( C S ) = E ∥ φ S ( S ) ∥ H S 2 ≤ κ S 2 . Moreover, C Z S C_{ZS} C Z S
∥ φ Z ( Z ) ⊗ φ S ( S ) ∥ H S = ∥ φ Z ( Z ) ∥ H Z ∥ φ S ( S ) ∥ H S ≤ κ Z κ S . \lVert\varphi_Z(Z)\otimes\varphi_S(S)\rVert_{\mathrm{HS}}=\lVert\varphi_Z(Z)\rVert_{\mathcal{H}_Z}\lVert\varphi_S(S)\rVert_{\mathcal{H}_S}\leq\kappa_Z\kappa_S. ∥ φ Z ( Z ) ⊗ φ S ( S ) ∥ HS = ∥ φ Z ( Z ) ∥ H Z ∥ φ S ( S ) ∥ H S ≤ κ Z κ S . Lemma D.4 (Consequences of attainability). Under Assumption D.2,
C Z S = H 0 C S . C_{ZS}=H_0C_S. C Z S = H 0 C S . Moreover, for every H ∈ S 2 ( H S , H Z ) H\in\mathcal{S}_2(\mathcal{H}_S,\mathcal{H}_Z) H ∈ S 2 ( H S , H Z )
∥ m H − m 0 ∥ L 2 ( P S , 0 ; H Z ) = ∥ ( H − H 0 ) C S 1 / 2 ∥ H S . \lVert m_H-m_0\rVert_{L_2(P_{S,0};\mathcal{H}_Z)}=\lVert(H-H_0)C_S^{1/2}\rVert_{\mathrm{HS}}. ∥ m H − m 0 ∥ L 2 ( P S , 0 ; H Z ) = ∥( H − H 0 ) C S 1/2 ∥ HS . Proof. For every f ∈ H S f\in\mathcal{H}_S f ∈ H S
C Z S f = E [ φ Z ( Z ) ⟨ φ S ( S ) , f ⟩ H S ] = E [ m 0 ( S ) ⟨ φ S ( S ) , f ⟩ H S ] = H 0 E [ φ S ( S ) ⟨ φ S ( S ) , f ⟩ H S ] = H 0 C S f . C_{ZS}f=\mathbb{E}\left[\varphi_Z(Z)\langle\varphi_S(S),f\rangle_{\mathcal{H}_S}\right]=\mathbb{E}\left[m_0(S)\langle\varphi_S(S),f\rangle_{\mathcal{H}_S}\right]=H_0\mathbb{E}\left[\varphi_S(S)\langle\varphi_S(S),f\rangle_{\mathcal{H}_S}\right]=H_0C_Sf. C Z S f = E [ φ Z ( Z ) ⟨ φ S ( S ) , f ⟩ H S ] = E [ m 0 ( S ) ⟨ φ S ( S ) , f ⟩ H S ] = H 0 E [ φ S ( S ) ⟨ φ S ( S ) , f ⟩ H S ] = H 0 C S f . This proves C Z S = H 0 C S C_{ZS}=H_0C_S C Z S = H 0 C S
∥ m H − m 0 ∥ L 2 ( P S , 0 ; H Z ) 2 = E [ ∥ ( H − H 0 ) φ S ( S ) ∥ H Z 2 ] = Tr [ ( H − H 0 ) C S ( H − H 0 ) ∗ ] = ∥ ( H − H 0 ) C S 1 / 2 ∥ H S 2 . \lVert m_H-m_0\rVert_{L_2(P_{S,0};\mathcal{H}_Z)}^2=\mathbb{E}\left[\lVert(H-H_0)\varphi_S(S)\rVert_{\mathcal{H}_Z}^2\right]=\operatorname{Tr}\left[(H-H_0)C_S(H-H_0)^*\right]=\lVert(H-H_0)C_S^{1/2}\rVert_{\mathrm{HS}}^2. ∥ m H − m 0 ∥ L 2 ( P S , 0 ; H Z ) 2 = E [ ∥( H − H 0 ) φ S ( S ) ∥ H Z 2 ] = Tr [ ( H − H 0 ) C S ( H − H 0 ) ∗ ] = ∥( H − H 0 ) C S 1/2 ∥ HS 2 . Taking nonnegative square roots completes the proof.
The next assumption characterizes the eigendecay of the input covariance operator and the maximal leverage of the input features. These quantities determine how many uniformly sampled landmarks are sufficient for the projected Nyström estimator to retain the statistical accuracy of full kernel ridge regression.
Assumption D.5 (Input spectral regularity; El Ahmad et al. [49 Assumption 3-4] ). For some γ ∈ ( 0 , 1 ] \gamma\in(0,1] γ ∈ ( 0 , 1 ]
Q S : = Tr ( C S γ ) < ∞ . Q_S:=\operatorname{Tr}(C_S^\gamma)<\infty. Q S := Tr ( C S γ ) < ∞. Moreover, for some μ ∈ [ 0 , 1 ) \mu\in[0,1) μ ∈ [ 0 , 1 ) B S < ∞ B_S<\infty B S < ∞ φ S ( s ) ⊗ φ S ( s ) ⪯ B S C S 1 − μ \varphi_S(s)\otimes\varphi_S(s)\preceq B_SC_S^{1-\mu} φ S ( s ) ⊗ φ S ( s ) ⪯ B S C S 1 − μ P S , 0 P_{S,0} P S , 0
Definition D.6 (Empirical sampling and covariance operators). Given observations ( S i , Z i ) i = 1 n (S_i,Z_i)_{i=1}^n ( S i , Z i ) i = 1 n
T S : H S → R n , ( T S f ) i : = 1 n f ( S i ) , T Z : H Z → R n , ( T Z g ) i : = 1 n g ( Z i ) . (51) T_S:\mathcal{H}_S\to\mathbb{R}^n,\qquad(T_Sf)_i:=\frac{1}{\sqrt{n}}f(S_i),\qquad T_Z:\mathcal{H}_Z\to\mathbb{R}^n,\qquad(T_Zg)_i:=\frac{1}{\sqrt{n}}g(Z_i). \tag*{(51)} T S : H S → R n , ( T S f ) i := n 1 f ( S i ) , T Z : H Z → R n , ( T Z g ) i := n 1 g ( Z i ) . (51) With R n \mathbb{R}^n R n
T S ∗ a = 1 n ∑ i = 1 n a i φ S ( S i ) , T Z ∗ a = 1 n ∑ i = 1 n a i φ Z ( Z i ) , a ∈ R n . (52) T_S^*a=\frac{1}{\sqrt{n}}\sum_{i=1}^n a_i\varphi_S(S_i),\qquad T_Z^*a=\frac{1}{\sqrt{n}}\sum_{i=1}^n a_i\varphi_Z(Z_i),\qquad a\in\mathbb{R}^n. \tag*{(52)} T S ∗ a = n 1 i = 1 ∑ n a i φ S ( S i ) , T Z ∗ a = n 1 i = 1 ∑ n a i φ Z ( Z i ) , a ∈ R n . (52) The empirical input covariance and output–input cross-covariance operators are respectively
C ^ S : = T S ∗ T S = 1 n ∑ i = 1 n φ S ( S i ) ⊗ φ S ( S i ) , C ^ Z S : = T Z ∗ T S = 1 n ∑ i = 1 n φ Z ( Z i ) ⊗ φ S ( S i ) . (53) \widehat{C}_S:=T_S^*T_S=\frac{1}{n}\sum_{i=1}^n\varphi_S(S_i)\otimes\varphi_S(S_i),\qquad\widehat{C}_{ZS}:=T_Z^*T_S=\frac{1}{n}\sum_{i=1}^n\varphi_Z(Z_i)\otimes\varphi_S(S_i). \tag*{(53)} C S := T S ∗ T S = n 1 i = 1 ∑ n φ S ( S i ) ⊗ φ S ( S i ) , C Z S := T Z ∗ T S = n 1 i = 1 ∑ n φ Z ( Z i ) ⊗ φ S ( S i ) . (53) Dense Hilbert-valued ridge regression. By Definition D.6, for every H ∈ S 2 ( H S , H Z ) H\in\mathcal{S}_2(\mathcal{H}_S,\mathcal{H}_Z) H ∈ S 2 ( H S , H Z )
∥ T Z ∗ − H T S ∗ ∥ H S 2 = 1 n ∑ i = 1 n ∥ φ Z ( Z i ) − H φ S ( S i ) ∥ H Z 2 . (54) \lVert T_Z^*-HT_S^*\rVert_{\mathrm{HS}}^2=\frac{1}{n}\sum_{i=1}^n\lVert\varphi_Z(Z_i)-H\varphi_S(S_i)\rVert_{\mathcal{H}_Z}^2. \tag*{(54)} ∥ T Z ∗ − H T S ∗ ∥ HS 2 = n 1 i = 1 ∑ n ∥ φ Z ( Z i ) − H φ S ( S i ) ∥ H Z 2 . (54) Consequently, for λ > 0 \lambda>0 λ > 0
H ^ λ ∈ arg min H ∈ S 2 ( H S , H Z ) { ∥ T Z ∗ − H T S ∗ ∥ H S 2 + λ ∥ H ∥ H S 2 } . (55) \widehat{H}_\lambda\in\underset{H\in\mathcal{S}_2(\mathcal{H}_S,\mathcal{H}_Z)}{\operatorname{arg\,min}}\left\{\lVert T_Z^*-HT_S^*\rVert_{\mathrm{HS}}^2+\lambda\lVert H\rVert_{\mathrm{HS}}^2\right\}. \tag*{(55)} H λ ∈ H ∈ S 2 ( H S , H Z ) arg min { ∥ T Z ∗ − H T S ∗ ∥ HS 2 + λ ∥ H ∥ HS 2 } . (55) Let K S : = ( k S ( S i , S j ) ) i , j = 1 n ∈ R n × n K_S:=(k_S(S_i,S_j))_{i,j=1}^n\in\mathbb{R}^{n\times n} K S := ( k S ( S i , S j ) ) i , j = 1 n ∈ R n × n s ∈ S s\in\mathcal{S} s ∈ S k S ( s ) : = ( k S ( S i , s ) ) i = 1 n ∈ R n k_S(s):=(k_S(S_i,s))_{i=1}^n\in\mathbb{R}^n k S ( s ) := ( k S ( S i , s ) ) i = 1 n ∈ R n K S = K A X K_S=K_{A X} K S = K A X
Proposition D.7 (Dense Hilbert-valued ridge solution; [50 Theorem (4.1)] , [49 Equation (3)–(4)] ). The optimization problem equation 51 admits a unique minimizer, given by
H ^ λ = C ^ Z S ( C ^ S + λ I H S ) − 1 . (56) \widehat{H}_{\lambda}=\widehat{C}_{ZS}\left(\widehat{C}_{S}+\lambda I_{\mathcal{H}_{S}}\right)^{-1}. \tag*{(56)} H λ = C Z S ( C S + λ I H S ) − 1 . (56) Consequently, writing m ^ λ ( s ) : = H ^ λ φ S ( s ) \widehat{m}_{\lambda}(s):=\widehat{H}_{\lambda}\varphi_{S}(s) m λ ( s ) := H λ φ S ( s )
m ^ λ ( s ) = ∑ i = 1 n [ ( K S + n λ I n ) − 1 k S ( s ) ] i φ Z ( Z i ) . (57) \widehat{m}_{\lambda}(s)=\sum_{i=1}^{n}\left[(K_{S}+n\lambda I_{n})^{-1}k_{S}(s)\right]_{i}\varphi_{Z}(Z_{i}). \tag*{(57)} m λ ( s ) = i = 1 ∑ n [ ( K S + nλ I n ) − 1 k S ( s ) ] i φ Z ( Z i ) . (57) Generic input sketch and Nyström specialization. Following El Ahmad et al. [49 ] , let R ∈ R m × n R\in\mathbb{R}^{m\times n} R ∈ R m × n [49 ] formulate their method for general input and output sketches; here we sketch only the input RKHS and leave the latent-output RKHS unchanged.
Define the input-sketch operator and its adjoint
B R : = R T S : H S → R m , B R ∗ = T S ⊤ R ⊤ : R m → H S , (58) B_{R}:=RT_{S}:\mathcal{H}_{S}\to\mathbb{R}^{m},\qquad B_{R}^{*}=T_{S}^{\top}R^{\top}:\mathbb{R}^{m}\to\mathcal{H}_{S}, \tag*{(58)} B R := R T S : H S → R m , B R ∗ = T S ⊤ R ⊤ : R m → H S , (58) The input-sketch operator induces the finite dimensional subspace
H S , R : = Range ( B R ∗ ) ⊆ H S . \mathcal{H}_{S,R}:=\operatorname{Range}(B_{R}^{*})\subseteq\mathcal{H}_{S}. H S , R := Range ( B R ∗ ) ⊆ H S . To make this explicit, let e j e_{j} e j j j j R m \mathbb{R}^{m} R m
B R ∗ e j = 1 n ∑ i = 1 n R j i φ S ( S i ) . B_{R}^{*}e_{j}=\frac{1}{\sqrt{n}}\sum_{i=1}^{n}R_{ji}\varphi_{S}(S_{i}). B R ∗ e j = n 1 i = 1 ∑ n R j i φ S ( S i ) . Consequently,
H S , R = span { ∑ i = 1 n R j i φ S ( S i ) : j = 1 , … , m } . \mathcal{H}_{S,R}=\operatorname{span}\left\{\sum_{i=1}^{n}R_{ji}\varphi_{S}(S_{i}):j=1,\ldots,m\right\}. H S , R = span { i = 1 ∑ n R j i φ S ( S i ) : j = 1 , … , m } . The orthogonal projector onto this subspace is
P R : = B R ∗ ( B R B R ∗ ) † B R , (59) P_{R}:=B_{R}^{*}(B_{R}B_{R}^{*})^{\dagger}B_{R}, \tag*{(59)} P R := B R ∗ ( B R B R ∗ ) † B R , (59) where † \dagger † r R = rank ( B R ) r_{R}=\operatorname{rank}(B_{R}) r R = rank ( B R )
B R = U R Σ R V R ∗ , B_{R}=U_{R}\Sigma_{R}V_{R}^{*}, B R = U R Σ R V R ∗ , where U R : R r R → R m U_{R}:\mathbb{R}^{r_{R}}\to\mathbb{R}^{m} U R : R r R → R m V R : R r R → H S V_{R}:\mathbb{R}^{r_{R}}\to\mathcal{H}_{S} V R : R r R → H S Σ R \Sigma_{R} Σ R
P R = V R V R ∗ . P_{R}=V_{R}V_{R}^{*}. P R = V R V R ∗ . For the Nyström estimator used here, choose 1 ≤ m ≤ n 1\leq m\leq n 1 ≤ m ≤ n
L = { ℓ 1 , … , ℓ m } ⊂ { 1 , … , n } L=\{\ell_{1},\ldots,\ell_{m}\}\subset\{1,\ldots,n\} L = { ℓ 1 , … , ℓ m } ⊂ { 1 , … , n } uniformly without replacement [31 ] , and define the coordinate-selection matrix R L ∈ R m × n R_{L}\in\mathbb{R}^{m\times n} R L ∈ R m × n a ∈ R n a\in\mathbb{R}^{n} a ∈ R n
( R L a ) j = a ℓ j , j = 1 , … , m . (R_{L}a)_{j}=a_{\ell_{j}},\qquad j=1,\ldots,m. ( R L a ) j = a ℓ j , j = 1 , … , m . Next, define the landmark synthesis operator as
Φ L : R m → H S , Φ L c = ∑ j = 1 m c j φ S ( S ℓ j ) . \Phi_{L}:\mathbb{R}^{m}\to\mathcal{H}_{S},\qquad\Phi_{L}c=\sum_{j=1}^{m}c_{j}\varphi_{S}(S_{\ell_{j}}). Φ L : R m → H S , Φ L c = j = 1 ∑ m c j φ S ( S ℓ j ) . Its adjoint is the landmark evaluation operator given by
Φ L ∗ f = [ f ( S ℓ 1 ) ⋮ f ( S ℓ m ) ] ∈ R m \Phi_{L}^{*}f=
\begin{bmatrix}
f(S_{\ell_{1}})\\
\vdots\\
f(S_{\ell_{m}})
\end{bmatrix}
\in\mathbb{R}^{m} Φ L ∗ f = f ( S ℓ 1 ) ⋮ f ( S ℓ m ) ∈ R m The normalized sampling operator and the landmark synthesis operator are related by
R L T S = 1 n Φ L ∗ , T S ∗ R L ⊤ = 1 n Φ L , R_L T_S = \frac{1}{\sqrt{n}}\Phi_L^*, \qquad T_S^*R_L^\top= \frac{1}{\sqrt{n}}\Phi_L, R L T S = n 1 Φ L ∗ , T S ∗ R L ⊤ = n 1 Φ L , Consequently, the subspace induced by the Nyström sketch is
H S , R L = Range ( T S ∗ R L ⊤ ) = Range ( Φ L ) = span { φ S ( S ℓ ) : ℓ ∈ L } . \mathcal{H}_{S,R_L} = \operatorname{Range}(T_S^*R_L^\top) = \operatorname{Range}(\Phi_L) = \operatorname{span}\{\varphi_S(S_\ell): \ell\in L\}. H S , R L = Range ( T S ∗ R L ⊤ ) = Range ( Φ L ) = span { φ S ( S ℓ ) : ℓ ∈ L } . Thus the Nyström approximation restricts the dependence of the conditional mean embedding to the span of the landmark input features. Define the landmark Gram matrix W : = Φ L ∗ Φ L W := \Phi_L^*\Phi_L W := Φ L ∗ Φ L
W j r = ⟨ φ S ( S ℓ j ) , φ ( S ℓ r ) ⟩ H S = k S ( S ℓ j , S ℓ r ) . W_{jr} = \langle\varphi_S(S_{\ell_j}), \varphi(S_{\ell_r})\rangle_{\mathcal{H}_S} = k_S(S_{\ell_j}, S_{\ell_r}). W j r = ⟨ φ S ( S ℓ j ) , φ ( S ℓ r ) ⟩ H S = k S ( S ℓ j , S ℓ r ) . Hence,
W = R L K S R L ⊤ = K S [ L , L ] ∈ R m × m . W = R_LK_SR_L^\top= K_S[L,L] \in\mathbb{R}^{m \times m}. W = R L K S R L ⊤ = K S [ L , L ] ∈ R m × m . The generic sketch projector now becomes
P L = P R L = Φ L W † Φ L ∗ . (60) P_L = P_{R_L} = \Phi_LW^\dagger\Phi_L^*. \tag*{(60)} P L = P R L = Φ L W † Φ L ∗ . (60) Projected regression problem. The sketch projector P R P_R P R
H S = H S , R ⊕ H S , R ⊥ , H S , R = Range ( P R ) . \mathcal{H}_S = \mathcal{H}_{S,R} \oplus\mathcal{H}_{S,R}^{\perp}, \qquad\mathcal{H}_{S,R} = \operatorname{Range}(P_R). H S = H S , R ⊕ H S , R ⊥ , H S , R = Range ( P R ) . We restrict the regression operator to use only the component of the input feature lying in H S , R \mathcal{H}_{S,R} H S , R
S 2 , R : = { H ∈ S ( H S , H Z ) : H = H P R } . \mathcal{S}_{2,R} := \{H \in\mathcal{S}(\mathcal{H}_S,\mathcal{H}_Z): H = HP_R\}. S 2 , R := { H ∈ S ( H S , H Z ) : H = H P R } . The condition H = H P R H = HP_R H = H P R H ( I H S − P R ) = 0 H(I_{\mathcal{H}_S} - P_R) = 0 H ( I H S − P R ) = 0 H H H
m H ( s ) = H φ S ( s ) = H P R φ S ( s ) . m_H(s) = H\varphi_S(s) = HP_R\varphi_S(s). m H ( s ) = H φ S ( s ) = H P R φ S ( s ) . Therefore, the estimator depends on s s s P R φ S ( s ) P_R\varphi_S(s) P R φ S ( s )
G R : = { m H : m H ( s ) = H φ S ( s ) , H ∈ S 2 , R } . \mathcal{G}_R := \{m_H : m_H(s) = H\varphi_S(s),\ H \in\mathcal{S}_{2,R}\}. G R := { m H : m H ( s ) = H φ S ( s ) , H ∈ S 2 , R } . Next, we define the input-sketched ridge estimator by
H ^ R , λ ∈ arg min H ∈ S 2 , R { ∥ T Z ∗ − H T S ∗ ∥ H S 2 + λ ∥ H ∥ H S 2 } . (61) \widehat{H}_{R,\lambda} \in\underset{H \in\mathcal{S}_{2,R}}{\arg\min}\left\{\lVert T_Z^* - HT_S^*\rVert_{\mathcal{H}_S}^2 + \lambda\lVert H\rVert_{\mathcal{H}_S}^2\right\}. \tag*{(61)} H R , λ ∈ H ∈ S 2 , R arg min { ∥ T Z ∗ − H T S ∗ ∥ H S 2 + λ ∥ H ∥ H S 2 } . (61) Equivalently,
H ^ R , λ ∈ arg min H ∈ S ( H S , H Z ) H = H P R { 1 n ∑ i = 1 n ∥ φ Z ( Z i ) − H P R φ S ( S i ) ∥ H Z 2 + λ ∥ H ∥ H S 2 } . (62) \widehat{H}_{R,\lambda} \in\underset{\substack{H \in\mathcal{S}(\mathcal{H}_S,\mathcal{H}_Z)\\ H=HP_R}}{\arg\min}\left\{\frac{1}{n}\sum_{i=1}^{n}\lVert\varphi_Z(Z_i)-HP_R\varphi_S(S_i)\rVert_{\mathcal{H}_Z}^2+\lambda\lVert H\rVert_{\mathcal{H}_S}^2\right\}. \tag*{(62)} H R , λ ∈ H ∈ S ( H S , H Z ) H = H P R arg min { n 1 i = 1 ∑ n ∥ φ Z ( Z i ) − H P R φ S ( S i ) ∥ H Z 2 + λ ∥ H ∥ H S 2 } . (62) It is ordinary Hilbert-valued ridge regression after replacing each input feature φ S ( S i ) \varphi_S(S_i) φ S ( S i ) [49 Section 3 and Appendix C, Lemma 1, Equations (29)–(34)] . For Nyström row selection, this is the Hilbert-valued analogue of the landmark-subspace estimator of Rudi et al. [31 Section 2, Eqs. (3)–(5)] .
Recall that r R = rank ( B R ) r_R = \operatorname{rank}(B_R) r R = rank ( B R ) B R = U R Σ R V R ∗ B_R = U_R\Sigma_RV_R^* B R = U R Σ R V R ∗ V R : R r R → H S V_R : \mathbb{R}^{r_R} \to\mathcal{H}_S V R : R r R → H S V R ∗ V R = I r R V_R^*V_R = I_{r_R} V R ∗ V R = I r R V R V R ∗ = P R V_RV_R^* = P_R V R V R ∗ = P R V R V_R V R H S , R = Range ( P R ) \mathcal{H}_{S,R} = \operatorname{Range}(P_R) H S , R = Range ( P R ) f ∈ H S , R f \in\mathcal{H}_{S,R} f ∈ H S , R
f = V R a , a = V R ∗ f ∈ R r R . f = V_Ra, \qquad a = V_R^*f \in\mathbb{R}^{r_R}. f = V R a , a = V R ∗ f ∈ R r R . We define the reduced empirical covariance matrix by
C ^ S , R r e d : = V R ∗ C ^ S V R ∈ R r R × r R . (63) \widehat{C}_{S,R}^{\mathrm{red}} := V_R^*\widehat{C}_SV_R \in\mathbb{R}^{r_R \times r_R}. \tag*{(63)} C S , R red := V R ∗ C S V R ∈ R r R × r R . (63) This matrix is the coordinate representation of the empirical covariance operator C ^ S \widehat{C}_S C S a ∈ R r a \in\mathbb{R}^r a ∈ R r
a ⊤ C ^ S , R r e d a = ⟨ V R a , C ^ S V R a ⟩ H S = 1 n ∑ i = 1 n ∣ ⟨ V R a , φ S ( S i ) ⟩ H S ∣ 2 ≥ 0. a^\top\widehat{C}^{\mathrm{red}}_{S,R}a = \left\langle V_Ra,\widehat{C}_SV_Ra\right\rangle_{\mathcal{H}_S} = \frac{1}{n}\sum_{i=1}^{n}\left|\left\langle V_Ra,\varphi_S(S_i)\right\rangle_{\mathcal{H}_S}\right|^2 \ge0. a ⊤ C S , R red a = ⟨ V R a , C S V R a ⟩ H S = n 1 i = 1 ∑ n ⟨ V R a , φ S ( S i ) ⟩ H S 2 ≥ 0. Thus C ^ S , R r e d \widehat{C}^{\mathrm{red}}_{S,R} C S , R red λ > 0 \lambda> 0 λ > 0 C ^ S , R r e d + λ I R ⪰ λ I R \widehat{C}^{\mathrm{red}}_{S,R}+\lambda I_R \succeq\lambda I_R C S , R red + λ I R ⪰ λ I R
Definition D.8 (Restricted regularized inverse). Following El Ahmad et al. [49, Lemma 1, Equation (29)], define the restricted regularized inverse for λ > 0 \lambda> 0 λ > 0
η R , λ ( C ^ S ) : = V R ( C ^ S , R r e d + λ I R ) − 1 V R ∗ : H S → H S . (64) \eta_{R,\lambda}(\widehat{C}_S) := V_R\left(\widehat{C}^{\mathrm{red}}_{S,R}+\lambda I_R\right)^{-1}V_R^*:\mathcal{H}_S\to\mathcal{H}_S. \tag*{(64)} η R , λ ( C S ) := V R ( C S , R red + λ I R ) − 1 V R ∗ : H S → H S . (64) The following result is the input-only specialization of El Ahmad et al. [49, Appendix C, Lemma 1, Eq. (29)], with the derivation in Eqs. (35) –(38) .
Proposition D.9 (Projected Hilbert-valued ridge solution). The optimization problem equation 57 admits a unique minimizer, given by
H ^ R , λ = C ^ Z S η R , λ ( C ^ S ) . (65) \widehat{H}_{R,\lambda}=\widehat{C}_{ZS}\eta_{R,\lambda}(\widehat{C}_S). \tag*{(65)} H R , λ = C Z S η R , λ ( C S ) . (65) Moreover, the restricted inverse satisfies
P R ( C ^ S + λ I H S ) η R , λ ( C ^ S ) = P R , (66) P_R\left(\widehat{C}_S+\lambda I_{\mathcal{H}_S}\right)\eta_{R,\lambda}(\widehat{C}_S)=P_R, \tag*{(66)} P R ( C S + λ I H S ) η R , λ ( C S ) = P R , (66) η R , λ ( C ^ S ) ( C ^ S + λ I H S ) P R = P R . (67) \eta_{R,\lambda}(\widehat{C}_S)\left(\widehat{C}_S+\lambda I_{\mathcal{H}_S}\right)P_R=P_R. \tag*{(67)} η R , λ ( C S ) ( C S + λ I H S ) P R = P R . (67) The purpose of η R , λ ( C ^ S ) \eta_{R,\lambda}(\widehat{C}_S) η R , λ ( C S )
Population ridge comparison. From this point onward, we specialize the generic sketch to the Nyström landmark sketch R L R_L R L
P L : = P R L , η L , λ : = η R L , λ ( C ^ S ) . P_L:=P_{R_L},\qquad\eta_{L,\lambda}:=\eta_{R_L,\lambda}(\widehat{C}_S). P L := P R L , η L , λ := η R L , λ ( C S ) . Recall that the corresponding projected estimator is
H ^ L , λ = C ^ Z S η L , λ , m ^ L , λ ( s ) : = H ^ L , λ φ S ( s ) . \widehat{H}_{L,\lambda}=\widehat{C}_{ZS}\eta_{L,\lambda},\qquad\widehat{m}_{L,\lambda}(s):=\widehat{H}_{L,\lambda}\varphi_S(s). H L , λ = C Z S η L , λ , m L , λ ( s ) := H L , λ φ S ( s ) . For λ > 0 \lambda> 0 λ > 0 C S , λ : = C S + λ I H S C_{S,\lambda}:=C_S+\lambda I_{\mathcal{H}_S} C S , λ := C S + λ I H S C S C_S C S C S , λ ⪰ λ I H S C_{S,\lambda}\succeq\lambda I_{\mathcal{H}_S} C S , λ ⪰ λ I H S ∥ C S , λ − 1 ∥ o p ≤ λ − 1 \left\|C_{S,\lambda}^{-1}\right\|_{\mathrm{op}}\le\lambda^{-1} C S , λ − 1 op ≤ λ − 1
Following the input-only specialization of El Ahmad et al. [49, Appendix D, proof of Theorem 1, Eqs. (43)–(48)], we introduce the following comparison operator and error decomposition.
Lemma D.10 (Projected Nyström error decomposition). Suppose Assumptions D.1 and D.2 hold. For λ > 0 \lambda> 0 λ > 0 C S , λ : = C S + λ I H S C_{S,\lambda}:=C_S+\lambda I_{\mathcal{H}_S} C S , λ := C S + λ I H S H λ : = H 0 C S C S , λ − 1 H_\lambda:=H_0C_SC_{S,\lambda}^{-1} H λ := H 0 C S C S , λ − 1
H λ − H 0 = − λ H 0 C S , λ − 1 . (68) H_\lambda-H_0=-\lambda H_0C_{S,\lambda}^{-1}. \tag*{(68)} H λ − H 0 = − λ H 0 C S , λ − 1 . (68) For every observed sample and every realized landmark set,
H ^ L , λ − H 0 = ( C ^ Z S − H λ C ^ S ) η L , λ + H λ ( C ^ S η L , λ − I H S ) + ( H λ − H 0 ) . (69) \widehat{H}_{L,\lambda}-H_0=\left(\widehat{C}_{ZS}-H_\lambda\widehat{C}_S\right)\eta_{L,\lambda}+H_\lambda\left(\widehat{C}_S\eta_{L,\lambda}-I_{\mathcal{H}_S}\right)+(H_\lambda-H_0). \tag*{(69)} H L , λ − H 0 = ( C Z S − H λ C S ) η L , λ + H λ ( C S η L , λ − I H S ) + ( H λ − H 0 ) . (69) Consequently,
r m ( L , λ ) : = ∥ m ^ L , λ − m 0 ∥ L 2 ( P S , 0 ; H Z ) ≤ ( A ) L , λ + ( B ) L , λ + ( C ) λ , (70) r_m(L,\lambda):=\left\|\widehat{m}_{L,\lambda}-m_0\right\|_{L^2(P_{S,0};\mathcal{H}_Z)}
\le(A)_{L,\lambda}+(B)_{L,\lambda}+(C)_\lambda, \tag*{(70)} r m ( L , λ ) := ∥ m L , λ − m 0 ∥ L 2 ( P S , 0 ; H Z ) ≤ ( A ) L , λ + ( B ) L , λ + ( C ) λ , (70) where
( A ) L , λ : = ∥ ( C ^ Z S − H λ C ^ S ) η L , λ C S 1 / 2 ∥ H S , (A)_{L,\lambda}:=\left\|\left(\widehat{C}_{ZS}-H_\lambda\widehat{C}_S\right)\eta_{L,\lambda}C_S^{1/2}\right\|_{\mathrm{HS}}, ( A ) L , λ := ( C Z S − H λ C S ) η L , λ C S 1/2 HS , ( B ) L , λ : = ∥ H Λ ( C ^ S η L , λ − I H S ) C S 1 / 2 ∥ H S , ( C ) λ : = ∥ ( H λ − H 0 ) C S 1 / 2 ∥ H S . (71) (B)_{L,\lambda} := \left\|H_{\Lambda}\left(\widehat{C}_{S\eta L,\lambda}-I_{\mathcal{H}_S}\right)C_S^{1/2}\right\|_{\mathcal{H}_S},\\
(C)_{\lambda} := \left\|(H_{\lambda}-H_0)C_S^{1/2}\right\|_{\mathcal{H}_S}. \tag*{(71)} ( B ) L , λ := H Λ ( C S η L , λ − I H S ) C S 1/2 H S , ( C ) λ := ( H λ − H 0 ) C S 1/2 H S . (71) Proof. Since C S C S , λ − 1 = I H S − λ C S , λ − 1 C_S C_{S,\lambda}^{-1}=I_{\mathcal{H}_S}-\lambda C_{S,\lambda}^{-1} C S C S , λ − 1 = I H S − λ C S , λ − 1 (68) follows immediately. Using H ^ L , λ = C ^ Z S η L , λ \widehat{H}_{L,\lambda}=\widehat{C}_{ZS}\eta_{L,\lambda} H L , λ = C Z S η L , λ H λ C ^ S η L , λ H_{\lambda}\widehat{C}_{S}\eta_{L,\lambda} H λ C S η L , λ H λ H_{\lambda} H λ (69) . The prediction-norm identity in Lemma D.4 , followed by the triangle inequality, yields Equation (70) . □
Lemma D.10 gives an exact algebraic decomposition of the projected Nyström estimation error into three components, which will be controlled separately. Because the decomposition holds pathwise for every realized sample and landmark set, it does not yet require the input-capacity or embedding assumptions, uniform landmark sampling, or concentration arguments. The sample term ( A ) L , λ (A)_{L,\lambda} ( A ) L , λ ( B ) L , λ (B)_{L,\lambda} ( B ) L , λ ( C ) λ (C)_{\lambda} ( C ) λ
Auxiliary operator bounds. For λ > 0 \lambda>0 λ > 0
C ^ S , λ : = C ^ S + λ I H S . \widehat{C}_{S,\lambda}:=\widehat{C}_S+\lambda I_{\mathcal{H}_S}. C S , λ := C S + λ I H S . We use the covariance-comparison event
E c m p ( λ ) : = { ∥ C ^ S , λ − 1 / 2 C S , λ 1 / 2 ∥ o p ≤ 2 , ∥ C S , λ − 1 / 2 C ^ S , λ 1 / 2 ∥ o p ≤ 3 2 } . (72) \mathcal{E}_{\mathrm{cmp}}(\lambda):=\left\{\left\|\widehat{C}_{S,\lambda}^{-1/2}C_{S,\lambda}^{1/2}\right\|_{\mathrm{op}}\le\sqrt{2},\ \left\|C_{S,\lambda}^{-1/2}\widehat{C}_{S,\lambda}^{1/2}\right\|_{\mathrm{op}}\le\sqrt{\frac{3}{2}}\right\}. \tag*{(72)} E cmp ( λ ) := { C S , λ − 1/2 C S , λ 1/2 op ≤ 2 , C S , λ − 1/2 C S , λ 1/2 op ≤ 2 3 } . (72) These are the input-space specializations of El Ahmad et al. [49, Lemmas 9–10, Eqs. (106)–(111)]. Applying the two lemmas with failure probability δ / 2 \delta/2 δ /2 E c m p ( λ ) \mathcal{E}_{\mathrm{cmp}}(\lambda) E cmp ( λ ) 1 − δ 1-\delta 1 − δ
We also use the deterministic restricted-resolvent contraction
∥ C ^ S , λ 1 / 2 η L , λ C ^ S , λ 1 / 2 ∥ o p ≤ 1 , (73) \left\|\widehat{C}_{S,\lambda}^{1/2}\eta_{L,\lambda}\widehat{C}_{S,\lambda}^{1/2}\right\|_{\mathrm{op}}\le1, \tag*{(73)} C S , λ 1/2 η L , λ C S , λ 1/2 op ≤ 1 , (73) which is the input-only specialization of El Ahmad et al. [49, Lemma 13, Eq. (124)]; see also Rudi et al. [31, Lemma 8].
Finally, define the input-subspace approximation error
ε L : = ∥ ( I H S − P L ) C S 1 / 2 ∥ o p . (74) \varepsilon_L:=\left\|(I_{\mathcal{H}_S}-P_L)C_S^{1/2}\right\|_{\mathrm{op}}. \tag*{(74)} ε L := ( I H S − P L ) C S 1/2 op . (74) The following bound is the input-only specialization of El Ahmad et al. [49, Lemma 3, Eq. (55)], with the proof in Eqs. (60) –(63) .
Lemma D.11 (Bound for the sample term). Suppose the observations are i.i.d., the input and latent kernels are bounded, Assumption D.2 and D.5 hold. Set
λ n : = n − 1 / ( 1 + γ ) . \lambda_n:=n^{-1/(1+\gamma)}. λ n := n − 1/ ( 1 + γ ) . Let δ ∈ ( 0 , 1 ] \delta\in(0,1] δ ∈ ( 0 , 1 ] n n n
9 κ S 2 n log ( n δ ) ≤ λ n ≤ ∥ C S ∥ o p . (75) \frac{9\kappa_S^2}{n}\log\left(\frac{n}{\delta}\right)\le\lambda_n\le\|C_S\|_{\mathrm{op}}. \tag*{(75)} n 9 κ S 2 log ( δ n ) ≤ λ n ≤ ∥ C S ∥ op . (75) Then there exists a finite problem-dependent constant C A > 0 C_A>0 C A > 0 L L L 1 − δ 1-\delta 1 − δ
( A ) L , λ n ≤ C A log ( 4 δ ) n − 1 2 ( 1 + γ ) . (76) (A)_{L,\lambda_n}\le C_A\log\left(\frac{4}{\delta}\right)n^{-\frac{1}{2(1+\gamma)}}. \tag*{(76)} ( A ) L , λ n ≤ C A log ( δ 4 ) n − 2 ( 1 + γ ) 1 . (76) Proof. Under the identifications X ≡ S X \equiv S X ≡ S Y ≡ Z Y \equiv Z Y ≡ Z η ~ ( C ^ X ) ≡ η L , λ ( C ^ S ) \widetilde{\eta}(\widehat{C}_X) \equiv\eta_{L,\lambda}(\widehat{C}_S) η ( C X ) ≡ η L , λ ( C S ) I H Z I_{\mathcal{H}_Z} I H Z ( A ) L , λ (A)_{L,\lambda} ( A ) L , λ ( A ) (A) ( A ) [49 Theorem 1] . The result therefore follows directly from El Ahmad et al. [49 Lemma 3, Eqs. (55)–(59)] .
The following result is the input-only specialization of El Ahmad et al. [49 Lemma 4, Eq. (60)] , with the proof in Eqs. (61)–(73)]. We use the sharper contraction ∥ H λ ∥ H S ≤ ∥ H 0 ∥ H S \lVert H_\lambda\rVert_{\mathrm{HS}} \le\lVert H_0\rVert_{\mathrm{HS}} ∥ H λ ∥ HS ≤ ∥ H 0 ∥ HS
Lemma D.12 (Bound for the projection/resolvent term). For every λ > 0 \lambda> 0 λ > 0 and every orthogonal input-sketch projector P L P_L P L on the covariance-comparison event E c m p ( λ ) \mathcal{E}_{\mathrm{cmp}}(\lambda) E cmp ( λ )
( B ) L , λ ≤ 2 3 ∥ H 0 ∥ H S ( ε L + λ ) . (77) (B)_{L,\lambda} \le2\sqrt{3}\lVert H_0\rVert_{\mathrm{HS}}\left(\varepsilon_L+\sqrt{\lambda}\right). \tag*{(77)} ( B ) L , λ ≤ 2 3 ∥ H 0 ∥ HS ( ε L + λ ) . (77) Proof. Set C ^ S , λ : = C ^ S + λ I H S \widehat{C}_{S,\lambda}:=\widehat{C}_S+\lambda I_{\mathcal{H}_S} C S , λ := C S + λ I H S P L C ^ S , λ η L , λ = P L P_L\widehat{C}_{S,\lambda}\eta_{L,\lambda}=P_L P L C S , λ η L , λ = P L
C ^ S η L , λ − I H S = ( I H S − P L ) C ^ S , λ η L , λ − λ η L , λ − ( I H S − P L ) . (78) \widehat{C}_S\eta_{L,\lambda}-I_{\mathcal{H}_S}=(I_{\mathcal{H}_S}-P_L)\widehat{C}_{S,\lambda}\eta_{L,\lambda}-\lambda\eta_{L,\lambda}-(I_{\mathcal{H}_S}-P_L). \tag*{(78)} C S η L , λ − I H S = ( I H S − P L ) C S , λ η L , λ − λ η L , λ − ( I H S − P L ) . (78) Consequently,
( B ) L , λ ≤ ∥ H λ ∥ H S ( B 1 + B 2 + B 3 ) , (B)_{L,\lambda}\le\lVert H_\lambda\rVert_{\mathrm{HS}}(B_1+B_2+B_3), ( B ) L , λ ≤ ∥ H λ ∥ HS ( B 1 + B 2 + B 3 ) , where
B 1 : = ∥ ( I H S − P L ) C ^ S , λ η L , λ C S 1 / 2 ∥ o p , B 2 : = λ ∥ η L , λ C S 1 / 2 ∥ o p , B 3 : = ∥ ( I H S − P L ) C S 1 / 2 ∥ o p = ε L . B_1:=\left\lVert(I_{\mathcal{H}_S}-P_L)\widehat{C}_{S,\lambda}\eta_{L,\lambda}C_S^{1/2}\right\rVert_{\mathrm{op}},\quad B_2:=\lambda\left\lVert\eta_{L,\lambda}C_S^{1/2}\right\rVert_{\mathrm{op}},\quad B_3:=\left\lVert(I_{\mathcal{H}_S}-P_L)C_S^{1/2}\right\rVert_{\mathrm{op}}=\varepsilon_L. B 1 := ( I H S − P L ) C S , λ η L , λ C S 1/2 op , B 2 := λ η L , λ C S 1/2 op , B 3 := ( I H S − P L ) C S 1/2 op = ε L . Using the covariance-comparison inequalities and the deterministic restricted-resolvent contraction, the operator factorizations in El Ahmad et al. [49 proof of Lemma 4, Eqs. (70)–(72)] give
B 1 ≤ 3 ( ε L + λ ) , B 2 ≤ 2 λ . B_1\le\sqrt{3}\left(\varepsilon_L+\sqrt{\lambda}\right),\qquad B_2\le\sqrt{2\lambda}. B 1 ≤ 3 ( ε L + λ ) , B 2 ≤ 2 λ . Therefore,
B 1 + B 2 + B 3 ≤ 3 ( ε L + λ ) + 2 λ + ε L ≤ 2 3 ( ε L + λ ) . B_1+B_2+B_3\le\sqrt{3}\left(\varepsilon_L+\sqrt{\lambda}\right)+\sqrt{2\lambda}+\varepsilon_L\le2\sqrt{3}\left(\varepsilon_L+\sqrt{\lambda}\right). B 1 + B 2 + B 3 ≤ 3 ( ε L + λ ) + 2 λ + ε L ≤ 2 3 ( ε L + λ ) . Finally,
H λ = H 0 C S ( C S + λ I H S ) − 1 , H_\lambda=H_0C_S(C_S+\lambda I_{\mathcal{H}_S})^{-1}, H λ = H 0 C S ( C S + λ I H S ) − 1 , and spectral calculus gives
∥ C S ( C S + λ I H S ) − 1 ∥ o p ≤ 1. \left\lVert C_S(C_S+\lambda I_{\mathcal{H}_S})^{-1}\right\rVert_{\mathrm{op}}\le1. C S ( C S + λ I H S ) − 1 op ≤ 1. Hence ∥ H λ ∥ H S ≤ ∥ H 0 ∥ H S \lVert H_\lambda\rVert_{\mathrm{HS}}\le\lVert H_0\rVert_{\mathrm{HS}} ∥ H λ ∥ HS ≤ ∥ H 0 ∥ HS
The following lemma is the input-only counterpart of El Ahmad et al. [49 Lemma 5, Eqs. (74)–(81)] . Because the output RKHS is not sketched, their output-projection term vanishes. We use a direct spectral calculation to obtain a slightly sharper constant.
Lemma D.13 (Population ridge bias). For every λ > 0 \lambda> 0 λ > 0
( C ) λ ≤ λ 2 ∥ H 0 ∥ H S . (79) (C)_\lambda\le\frac{\sqrt{\lambda}}{2}\lVert H_0\rVert_{\mathrm{HS}}. \tag*{(79)} ( C ) λ ≤ 2 λ ∥ H 0 ∥ HS . (79) Proof. By equation 64, H λ − H 0 = − λ H 0 C S , λ − 1 H_\lambda-H_0=-\lambda H_0C_{S,\lambda}^{-1} H λ − H 0 = − λ H 0 C S , λ − 1
( C ) λ = λ ∥ H 0 C S , λ − 1 C S 1 / 2 ∥ H S ≤ ∥ H 0 ∥ H S ∥ λ C S , λ − 1 C S 1 / 2 ∥ o p . (C)_\lambda=\lambda\left\lVert H_0C_{S,\lambda}^{-1}C_S^{1/2}\right\rVert_{\mathrm{HS}}\le\lVert H_0\rVert_{\mathrm{HS}}\left\lVert\lambda C_{S,\lambda}^{-1}C_S^{1/2}\right\rVert_{\mathrm{op}}. ( C ) λ = λ H 0 C S , λ − 1 C S 1/2 HS ≤ ∥ H 0 ∥ HS λ C S , λ − 1 C S 1/2 op . Notice that
∥ λ C S , λ − 1 C S 1 / 2 ∥ o p = sup σ ∈ spec ( C S ) λ σ σ + λ . \left\lVert\lambda C_{S,\lambda}^{-1}C_S^{1/2}\right\rVert_{\mathrm{op}}=\sup_{\sigma\in\operatorname{spec}(C_S)}\frac{\lambda\sqrt{\sigma}}{\sigma+\lambda}. λ C S , λ − 1 C S 1/2 op = σ ∈ spec ( C S ) sup σ + λ λ σ . Since σ + λ ≥ 2 σ λ \sigma+\lambda\ge2\sqrt{\sigma\lambda} σ + λ ≥ 2 σ λ λ σ σ + λ ≤ λ 2 \frac{\lambda\sqrt{\sigma}}{\sigma+\lambda}\le\frac{\sqrt{\lambda}}{2} σ + λ λ σ ≤ 2 λ σ ≥ 0 \sigma\ge0 σ ≥ 0 ( C ) λ ≤ λ 2 ∥ H 0 ∥ H S (C)_\lambda\le\frac{\sqrt{\lambda}}{2}\lVert H_0\rVert_{\mathrm{HS}} ( C ) λ ≤ 2 λ ∥ H 0 ∥ HS
Reduction to the Nyström projection error. Combining Lemmas D.11, D.12, and D.13, we obtain, on E c m p ( λ ) \mathcal{E}_{\mathrm{cmp}}(\lambda) E cmp ( λ )
r m ( L , λ ) ≤ ( A ) L , λ + 2 3 ∥ H 0 ∥ H S ε L + ( 2 3 + 1 2 ) ∥ H 0 ∥ H S λ . (80) r_m(L,\lambda) \le(A)_{L,\lambda}+2\sqrt{3}\lVert H_0\rVert_{\mathrm{HS}}\,\varepsilon_L+\left(2\sqrt{3}+\frac{1}{2}\right)\lVert H_0\rVert_{\mathrm{HS}}\sqrt{\lambda}. \tag*{(80)} r m ( L , λ ) ≤ ( A ) L , λ + 2 3 ∥ H 0 ∥ HS ε L + ( 2 3 + 2 1 ) ∥ H 0 ∥ HS λ . (80) It remains to control the random subspace error ε L \varepsilon_L ε L
For λ > 0 \lambda>0 λ > 0
N ( λ ) : = Tr [ C S C S , λ − 1 ] , (81) \mathcal{N}(\lambda):=\operatorname{Tr}\left[C_S C_{S,\lambda}^{-1}\right], \tag*{(81)} N ( λ ) := Tr [ C S C S , λ − 1 ] , (81) and the maximal ridge leverage
N ∞ ( λ ) : = ess sup s ∼ P S , 0 ⟨ φ S ( s ) , C S , λ − 1 φ S ( s ) ⟩ H S . (82) \mathcal{N}_{\infty}(\lambda):=\underset{s\sim P_{S,0}}{\operatorname{ess\,sup}}\left\langle\varphi_S(s),C_{S,\lambda}^{-1}\varphi_S(s)\right\rangle_{\mathcal{H}_S}. \tag*{(82)} N ∞ ( λ ) := s ∼ P S , 0 ess sup ⟨ φ S ( s ) , C S , λ − 1 φ S ( s ) ⟩ H S . (82) The effective dimension controls the average regularized complexity of the input RKHS, whereas N ∞ ( λ ) \mathcal{N}_{\infty}(\lambda) N ∞ ( λ )
Lemma D.14 (Uniform Nyström projection control). Let L ⊂ { 1 , … , n } L\subset\{1,\ldots,n\} L ⊂ { 1 , … , n } m m m P L P_L P L
span { φ S ( S ℓ ) : ℓ ∈ L } . \operatorname{span}\{\varphi_S(S_\ell):\ell\in L\}. span { φ S ( S ℓ ) : ℓ ∈ L } . If
m ≥ { 67 ∨ 5 N ∞ ( λ ) } log ( 4 κ S 2 λ δ ) , (83) m\geq\{67\vee5\mathcal{N}_{\infty}(\lambda)\}\log\left(\frac{4\kappa_S^2}{\lambda\delta}\right), \tag*{(83)} m ≥ { 67 ∨ 5 N ∞ ( λ )} log ( λ δ 4 κ S 2 ) , (83) then, with probability at least 1 − δ 1-\delta 1 − δ
∥ ( I H S − P L ) C S , λ 1 / 2 ∥ o p 2 ≤ 3 λ . (84) \left\lVert(I_{\mathcal{H}_S}-P_L)C_{S,\lambda}^{1/2}\right\rVert_{\mathrm{op}}^2\leq3\lambda. \tag*{(84)} ( I H S − P L ) C S , λ 1/2 op 2 ≤ 3 λ . (84) Consequently,
ε L ≤ 3 λ . (85) \varepsilon_L\leq\sqrt{3\lambda}. \tag*{(85)} ε L ≤ 3 λ . (85) Proof. The regularized projection bound equation 80 is the plain Nyström approximation result of Rudi et al. [31 Lemma 6] , applied to the covariance operator C S C_S C S N ∞ ( λ ) \mathcal{N}_{\infty}(\lambda) N ∞ ( λ ) P S , 0 P_{S,0} P S , 0
For the second claim, write
C S 1 / 2 = C S , λ 1 / 2 C S , λ − 1 / 2 C S 1 / 2 . C_S^{1/2}=C_{S,\lambda}^{1/2}C_{S,\lambda}^{-1/2}C_S^{1/2}. C S 1/2 = C S , λ 1/2 C S , λ − 1/2 C S 1/2 . Hence
ε L ≤ ∥ ( I H S − P L ) C S , λ 1 / 2 ∥ o p ∥ C S , λ − 1 / 2 C S 1 / 2 ∥ o p . \varepsilon_L\leq\left\lVert(I_{\mathcal{H}_S}-P_L)C_{S,\lambda}^{1/2}\right\rVert_{\mathrm{op}}\left\lVert C_{S,\lambda}^{-1/2}C_S^{1/2}\right\rVert_{\mathrm{op}}. ε L ≤ ( I H S − P L ) C S , λ 1/2 op C S , λ − 1/2 C S 1/2 op . Spectral calculus gives
∥ C S , λ − 1 / 2 C S 1 / 2 ∥ o p = sup σ ∈ spec ( C S ) σ σ + λ ≤ 1. \left\lVert C_{S,\lambda}^{-1/2}C_S^{1/2}\right\rVert_{\mathrm{op}}=\sup_{\sigma\in\operatorname{spec}(C_S)}\sqrt{\frac{\sigma}{\sigma+\lambda}}\leq1. C S , λ − 1/2 C S 1/2 op = σ ∈ spec ( C S ) sup σ + λ σ ≤ 1. The conclusion follows from equation 80. Consequently, on the intersection of E c m p ( λ ) \mathcal{E}_{\mathrm{cmp}}(\lambda) E cmp ( λ )
r m ( L , λ ) ≤ ( A ) L , λ + C H λ , (86) r_m(L,\lambda)\leq(A)_{L,\lambda}+C_H\sqrt{\lambda}, \tag*{(86)} r m ( L , λ ) ≤ ( A ) L , λ + C H λ , (86) where
C H : = [ 2 3 ( 3 + 1 ) + 1 2 ] ∥ H 0 ∥ H S . C_H:=\left[2\sqrt{3}(\sqrt{3}+1)+\frac{1}{2}\right]\lVert H_0\rVert_{\mathrm{HS}}. C H := [ 2 3 ( 3 + 1 ) + 2 1 ] ∥ H 0 ∥ HS . Lemma D.15 (Capacity bounds for the leverage quantity). Under the capacity condition in Assumption D.5 , we have, for every λ > 0 \lambda> 0 λ > 0
N ( λ ) ≤ Q S λ − γ . (87) \mathcal{N}(\lambda) \le Q_{S}\lambda^{-\gamma}. \tag*{(87)} N ( λ ) ≤ Q S λ − γ . (87) Moreover, boundedness of $\kappa_{S} always gives
N ∞ ( λ ) ≤ κ S 2 λ . (88) \mathcal{N}_{\infty}(\lambda) \le\frac{\kappa_{S}^{2}}{\lambda}. \tag*{(88)} N ∞ ( λ ) ≤ λ κ S 2 . (88) If, in addition, the embedding condition holds and γ + μ ≤ 1 \gamma+\mu\le1 γ + μ ≤ 1 then
N ∞ ( λ ) ≤ B S Q S λ − ( γ + μ ) . (89) \mathcal{N}_{\infty}(\lambda) \le B_{S}Q_{S}\lambda^{-(\gamma+\mu)}. \tag*{(89)} N ∞ ( λ ) ≤ B S Q S λ − ( γ + μ ) . (89) Proof. The effective-dimension bound is the input-side specialization of El Ahmad et al. [49 Lemma 11, Eqs. (115)–(118)] . The universal bound follows from C S , λ − 1 ⪯ λ − 1 I C_{S,\lambda}^{-1} \preceq\lambda^{-1}I C S , λ − 1 ⪯ λ − 1 I ∥ φ S ( s ) ∥ H S 2 ≤ κ S 2 \|\varphi_{S}(s)\|_{\mathcal{H}_{S}}^{2} \le\kappa_{S}^{2} ∥ φ S ( s ) ∥ H S 2 ≤ κ S 2
For the final claim, the embedding condition gives
N ∞ ( λ ) ≤ B S ∑ j σ j 1 − μ σ j + λ . \mathcal{N}_{\infty}(\lambda) \le B_{S}\sum_{j}\frac{\sigma_{j}^{1-\mu}}{\sigma_{j}+\lambda}. N ∞ ( λ ) ≤ B S j ∑ σ j + λ σ j 1 − μ . where ( σ j ) j (\sigma_{j})_{j} ( σ j ) j C S C_{S} C S γ + μ ≤ 1 \gamma+\mu\le1 γ + μ ≤ 1
σ 1 − μ σ + λ ≤ λ − ( γ + μ ) σ γ \frac{\sigma^{1-\mu}}{\sigma+\lambda} \le\lambda^{-(\gamma+\mu)}\sigma^{\gamma} σ + λ σ 1 − μ ≤ λ − ( γ + μ ) σ γ for every σ ≥ 0 \sigma\ge0 σ ≥ 0 j j j
Theorem D.16 (Projected Nyström CME rate). Suppose Assumptions D.1 and D.2 hold, and suppose that, for some γ ∈ ( 0 , 1 ] \gamma\in(0,1] γ ∈ ( 0 , 1 ] Tr ( C S γ ) < ∞ \operatorname{Tr}(C_{S}^{\gamma}) < \infty Tr ( C S γ ) < ∞ Set
λ n : = n − 1 / ( 1 + γ ) . \lambda_{n} := n^{-1/(1+\gamma)}. λ n := n − 1/ ( 1 + γ ) . For each n n n let L n ⊆ { 1 , … , n } L_{n} \subseteq\{1,\ldots,n\} L n ⊆ { 1 , … , n } contain 1 ≤ m n ≤ n 1 \le m_{n} \le n 1 ≤ m n ≤ n indices sampled uniformly without replacement. Fix δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 ) Suppose n n n is sufficiently large that
9 κ S 2 n log ( 2 n δ ) ≤ λ n ≤ ∥ C S ∥ o p , (90) \frac{9\kappa_{S}^{2}}{n}\log\left(\frac{2n}{\delta}\right) \le\lambda_{n} \le\|C_{S}\|_{\mathrm{op}}, \tag*{(90)} n 9 κ S 2 log ( δ 2 n ) ≤ λ n ≤ ∥ C S ∥ op , (90) and suppose
m n ≥ { 67 ∨ 5 N ∞ ( λ n ) } log ( 8 κ S 2 λ n δ ) . (91) m_{n} \ge\{67 \vee5\mathcal{N}_{\infty}(\lambda_{n})\}\log\left(\frac{8\kappa_{S}^{2}}{\lambda_{n}\delta}\right). \tag*{(91)} m n ≥ { 67 ∨ 5 N ∞ ( λ n )} log ( λ n δ 8 κ S 2 ) . (91) Then there exists a finite problem-dependent constant C A > 0 C_{A} > 0 C A > 0 independent of n n n m n m_{n} m n and δ \delta δ such that, with probability at least 1 − δ 1-\delta 1 − δ
∥ m ^ L n , λ n − m 0 ∥ L 2 ( P S , 0 ; H Z ) ≤ { C A log ( 8 δ ) + C H } n − 1 2 ( 1 + γ ) . (92) \|\widehat{m}_{L_{n},\lambda_{n}}-m_{0}\|_{L_{2}(P_{S,0};\mathcal{H}_{Z})}
\le\left\{C_{A}\log\left(\frac{8}{\delta}\right)+C_{H}\right\}n^{-\frac{1}{2(1+\gamma)}}. \tag*{(92)} ∥ m L n , λ n − m 0 ∥ L 2 ( P S , 0 ; H Z ) ≤ { C A log ( δ 8 ) + C H } n − 2 ( 1 + γ ) 1 . (92) Proof. Apply Lemma D.11 with confidence level δ / 2 \delta/2 δ /2 1 − δ / 2 1-\delta/2 1 − δ /2
( A ) L n , λ n ≤ C A log ( 8 δ ) n − 1 2 ( 1 + γ ) , (A)_{L_{n},\lambda_{n}} \le C_{A}\log\left(\frac{8}{\delta}\right)n^{-\frac{1}{2(1+\gamma)}}, ( A ) L n , λ n ≤ C A log ( δ 8 ) n − 2 ( 1 + γ ) 1 , and the covariance-comparison event required by Lemma D.12 holds.
Condition in Equation (91) and Lemma D.14, applied with confidence level δ / 2 \delta/2 δ /2 1 − δ / 2 1-\delta/2 1 − δ /2
ε L n ≤ 3 λ n . \varepsilon_{L_{n}} \le\sqrt{3\lambda_{n}}. ε L n ≤ 3 λ n . On the intersection of these two events,
( B ) L n , λ n ≤ 2 3 ∥ H 0 ∥ H S ( ε L n + λ n ) ≤ 2 3 ( 3 + 1 ) ∥ H 0 ∥ H S λ n , (B)_{L_{n},\lambda_{n}} \le2\sqrt{3}\|H_{0}\|_{\mathrm{HS}}\left(\varepsilon_{L_{n}}+\sqrt{\lambda_{n}}\right) \le2\sqrt{3}(\sqrt{3}+1)\|H_{0}\|_{\mathrm{HS}}\sqrt{\lambda_{n}}, ( B ) L n , λ n ≤ 2 3 ∥ H 0 ∥ HS ( ε L n + λ n ) ≤ 2 3 ( 3 + 1 ) ∥ H 0 ∥ HS λ n , while Lemma D.13 gives
( C ) λ n ≤ 1 2 ∥ H 0 ∥ H S λ n . (C)_{\lambda_n} \le\frac{1}{2}\|H_0\|_{\mathrm{HS}}\sqrt{\lambda_n}. ( C ) λ n ≤ 2 1 ∥ H 0 ∥ HS λ n . Combining the three terms and using
λ n = n − 1 2 ( 1 + γ ) \sqrt{\lambda_n}=n^{-\frac{1}{2(1+\gamma)}} λ n = n − 2 ( 1 + γ ) 1 proves the claim.
Corollary D.17 (Explicit landmark schedule). Suppose Assumption D.5 holds, and define
ϑ : = min { γ + μ , 1 } , C ⋆ : = { B S Q S , γ + μ ≤ 1 , κ S 2 , γ + μ > 1. (93) \vartheta:= \min\{\gamma+\mu,1\}, \qquad C_{\star}:=
\begin{cases}
B_SQ_S, & \gamma+\mu\le1,\\
\kappa_S^2, & \gamma+\mu>1.
\end{cases} \tag*{(93)} ϑ := min { γ + μ , 1 } , C ⋆ := { B S Q S , κ S 2 , γ + μ ≤ 1 , γ + μ > 1. (93) Then
N ∞ ( λ ) ≤ C ⋆ λ − ϑ . (94) N_{\infty}(\lambda)\le C_{\star}\lambda^{-\vartheta}. \tag*{(94)} N ∞ ( λ ) ≤ C ⋆ λ − ϑ . (94) Consequently, for
λ n = n − 1 / ( 1 + γ ) , \lambda_n=n^{-1/(1+\gamma)}, λ n = n − 1/ ( 1 + γ ) , a sufficient finite-sample landmark condition is
m n ≥ { 67 ∨ 5 C ⋆ n ϑ 1 + γ } log ( 8 κ S 2 δ n 1 1 + γ ) . (95) m_n\ge\left\{67\vee5C_{\star}n^{\frac{\vartheta}{1+\gamma}}\right\}\log\left(\frac{8\kappa_S^2}{\delta}n^{\frac{1}{1+\gamma}}\right). \tag*{(95)} m n ≥ { 67 ∨ 5 C ⋆ n 1 + γ ϑ } log ( δ 8 κ S 2 n 1 + γ 1 ) . (95) For fixed δ \delta δ
m n ≳ n ϑ 1 + γ log n . (96) m_n\gtrsim n^{\frac{\vartheta}{1+\gamma}}\log n. \tag*{(96)} m n ≳ n 1 + γ ϑ log n . (96) Under this schedule,
∥ m ^ L n , λ n − m 0 ∥ L 2 ( P S , 0 ; H Z ) = O P ( n − 1 2 ( 1 + γ ) ) . (97) \|\widehat{m}_{L_n,\lambda_n}-m_0\|_{L_2(P_{S,0};\mathcal{H}_Z)}
=O_{\mathbb{P}}\left(n^{-\frac{1}{2(1+\gamma)}}\right). \tag*{(97)} ∥ m L n , λ n − m 0 ∥ L 2 ( P S , 0 ; H Z ) = O P ( n − 2 ( 1 + γ ) 1 ) . (97) Proof. The piecewise leverage bound follows from Lemma D.15: use equation 85 when γ + μ ≤ 1 \gamma+\mu\le1 γ + μ ≤ 1
λ n = n − 1 / ( 1 + γ ) \lambda_n=n^{-1/(1+\gamma)} λ n = n − 1/ ( 1 + γ ) into equation 87 gives equation 91. The rate then follows from Theorem D.16.
Same-sample embedding estimation. We next verify the outcome-side class control needed when the regression and the DR score use the same observations. All asymptotic statements below use the same ridge and landmark sequences, with the finite-sample conditions of Theorem D.16 holding eventually for each fixed δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 )
Lemma D.18 (Tightness of the projected regression operator). Under the conditions of Theorem D.16,
∥ H ^ L n , λ n ∥ H S = O P ( 1 ) , ∥ H ^ L n , λ n − H 0 ∥ o p = O P ( 1 ) . (98) \|\widehat{H}_{L_n,\lambda_n}\|_{\mathrm{HS}}=O_{\mathbb{P}}(1), \qquad
\|\widehat{H}_{L_n,\lambda_n}-H_0\|_{\mathrm{op}}=O_{\mathbb{P}}(1). \tag*{(98)} ∥ H L n , λ n ∥ HS = O P ( 1 ) , ∥ H L n , λ n − H 0 ∥ op = O P ( 1 ) . (98) Proof. Write T = C S T=C_S T = C S T ^ = C ^ S \widehat{T}=\widehat{C}_S T = C S P = P L n P=P_{L_n} P = P L n η = η L n , λ n \eta=\eta_{L_n,\lambda_n} η = η L n , λ n λ = λ n \lambda=\lambda_n λ = λ n
r 0 ( O ) = φ Z ( Z ) − H 0 φ S ( S ) , E n = P n [ r 0 ⊗ φ S ] . r_0(O)=\varphi_Z(Z)-H_0\varphi_S(S), \qquad E_n=P_n[r_0\otimes\varphi_S]. r 0 ( O ) = φ Z ( Z ) − H 0 φ S ( S ) , E n = P n [ r 0 ⊗ φ S ] . Then E [ r 0 ∣ S ] = 0 \mathbb{E}[r_0\mid S]=0 E [ r 0 ∣ S ] = 0 ∥ r 0 ∥ ≤ 2 κ Z \|r_0\|\le2\kappa_Z ∥ r 0 ∥ ≤ 2 κ Z H ^ L n , λ n = ( H 0 T ^ + E n ) η \widehat{H}_{L_n,\lambda_n}=(H_0\widehat{T}+E_n)\eta H L n , λ n = ( H 0 T + E n ) η
∥ ( I − P ) T ^ 1 / 2 ∥ o p ≤ 3 λ / 2 . \|(I-P)\widehat{T}^{1/2}\|_{\mathrm{op}}\le3\sqrt{\lambda/2}. ∥ ( I − P ) T 1/2 ∥ op ≤ 3 λ /2 . The restricted inverse identities give η ( T ^ + λ I ) η = η \eta(\widehat{T}+\lambda I)\eta=\eta η ( T + λ I ) η = η 0 ⪯ η ⪯ λ − 1 P 0\preceq\eta\preceq\lambda^{-1}P 0 ⪯ η ⪯ λ − 1 P
∥ T ^ 1 / 2 η ∥ o p ≤ 1 2 λ , ∥ ( T + λ I ) 1 / 2 η ∥ o p ≤ 2 / λ . \|\widehat{T}^{1/2}\eta\|_{\mathrm{op}}\le\frac{1}{2\sqrt{\lambda}}, \qquad
\|(T+\lambda I)^{1/2}\eta\|_{\mathrm{op}}\le\sqrt{2/\lambda}. ∥ T 1/2 η ∥ op ≤ 2 λ 1 , ∥ ( T + λ I ) 1/2 η ∥ op ≤ 2/ λ . Since T ^ η = P − λ η + ( I − P ) T ^ η \widehat{T}\eta= P-\lambda\eta+(I-P)\widehat{T}\eta T η = P − λ η + ( I − P ) T η ∥ T ^ η ∥ o p ≤ 1 + 3 / ( 2 2 ) \|\widehat{T}\eta\|_{\mathrm{op}}\le1+3/(2\sqrt{2}) ∥ T η ∥ op ≤ 1 + 3/ ( 2 2 )
E ∥ E n ( T + λ I ) − 1 / 2 ∥ H S 2 ≤ 4 κ Z 2 n N ( λ ) . \mathbb{E}\left\|E_n(T+\lambda I)^{-1/2}\right\|_{\mathrm{HS}}^2\le\frac{4\kappa_Z^2}{n}\mathcal{N}(\lambda). E E n ( T + λ I ) − 1/2 HS 2 ≤ n 4 κ Z 2 N ( λ ) . Factoring E n η E_n\eta E n η ( T + λ I ) − 1 / 2 (T+\lambda I)^{-1/2} ( T + λ I ) − 1/2
∥ H ^ L n , λ n ∥ H S = O P ( 1 + N ( λ n ) n λ n ) = O P ( 1 ) , \left\|\widehat{H}L_{n,\lambda_n}\right\|_{\mathrm{HS}}=O_{\mathbb{P}}\left(1+\sqrt{\frac{\mathcal{N}(\lambda_n)}{n\lambda_n}}\right)=O_{\mathbb{P}}(1), H L n , λ n HS = O P 1 + n λ n N ( λ n ) = O P ( 1 ) , because N ( λ ) ≤ Q S λ − γ \mathcal{N}(\lambda)\le Q_S\lambda^{-\gamma} N ( λ ) ≤ Q S λ − γ Q S = Tr ( C S γ ) < ∞ Q_S=\operatorname{Tr}(C_S^\gamma)<\infty Q S = Tr ( C S γ ) < ∞ n λ n 1 + γ = 1 n\lambda_n^{1+\gamma}=1 n λ n 1 + γ = 1 ∥ H ^ − H 0 ∥ o p ≤ ∥ H ^ ∥ H S + ∥ H 0 ∥ H S \|\widehat{H}-H_0\|_{\mathrm{op}}\le\|\widehat{H}\|_{\mathrm{HS}}+\|H_0\|_{\mathrm{HS}} ∥ H − H 0 ∥ op ≤ ∥ H ∥ HS + ∥ H 0 ∥ HS
For the exact-policy-integral estimator h ^ π = P n Γ η ^ \widehat{h}_\pi=P_n\Gamma_{\widehat{\eta}} h π = P n Γ η Γ 0 = Γ m 0 , w 0 \Gamma_0=\Gamma_{m_0,w_0} Γ 0 = Γ m 0 , w 0 G n = n ( P n − P 0 ) G_n=\sqrt{n}(P_n-P_0) G n = n ( P n − P 0 ) Γ 0 − χ ( π ) \Gamma_0-\chi(\pi) Γ 0 − χ ( π )
h ^ π − χ ( π ) = ( P n − P 0 ) Γ 0 + ( P n − P 0 ) ( Γ η ^ − Γ 0 ) − P 0 [ ( w ^ − w 0 ) ( m ^ − m 0 ) ] . (99) \widehat{h}_\pi-\chi(\pi)=(P_n-P_0)\Gamma_0+(P_n-P_0)(\Gamma_{\widehat{\eta}}-\Gamma_0)-P_0[(\widehat{w}-w_0)(\widehat{m}-m_0)]. \tag*{(99)} h π − χ ( π ) = ( P n − P 0 ) Γ 0 + ( P n − P 0 ) ( Γ η − Γ 0 ) − P 0 [( w − w 0 ) ( m − m 0 )] . (99) Here P 0 P_0 P 0 ∥ w ^ − w 0 ∥ L 2 ( P S , 0 ) ∥ m ^ − m 0 ∥ L 2 ( P S , 0 ; H Z ) \|\widehat{w}-w_0\|_{L_2(P_{S,0})}\|\widehat{m}-m_0\|_{L_2(P_{S,0};\mathcal{H}_Z)} ∥ w − w 0 ∥ L 2 ( P S , 0 ) ∥ m − m 0 ∥ L 2 ( P S , 0 ; H Z )
Proposition D.19 (Same-sample score control and equicontinuity). *Under Assumptions 2.1, D.1, and D.2, fix the target policy, encoder, and kernels, and suppose P 0 w 0 2 < ∞ P_0w_0^2<\infty P 0 w 0 2 < ∞ m ^ = H ^ φ S \widehat{m}=\widehat{H}\varphi_S m = H φ S ∥ H ^ − H 0 ∥ o p = O P ( 1 ) \|\widehat{H}-H_0\|_{\mathrm{op}}=O_{\mathbb{P}}(1) ∥ H − H 0 ∥ op = O P ( 1 ) q n = w ^ − w 0 q_n=\widehat{w}-w_0 q n = w − w 0 K w \mathcal{K}_w K w t t t P 0 ∥ t ( S ) ∥ K w 2 < ∞ P_0\|t(S)\|_{\mathcal{K}_w}^2<\infty P 0 ∥ t ( S ) ∥ K w 2 < ∞ ∥ q n ∥ K w = O P ( 1 ) \|q_n\|_{\mathcal{K}_w}=O_{\mathbb{P}}(1) ∥ q n ∥ K w = O P ( 1 )
∥ G n ( Γ η ^ − Γ 0 ) ∥ H Z = O P ( 1 ) . \|G_n(\Gamma_{\widehat{\eta}}-\Gamma_0)\|_{\mathcal{H}_Z}=O_{\mathbb{P}}(1). ∥ G n ( Γ η − Γ 0 ) ∥ H Z = O P ( 1 ) . If, additionally, both nuisance errors tend to zero in their respective L 2 ( P S , 0 ) L_2(P_{S,0}) L 2 ( P S , 0 ) o P ( 1 ) o_{\mathbb{P}}(1) o P ( 1 )
h ^ π − χ ( π ) = ( P n − P 0 ) Γ 0 − P 0 [ ( w ^ − w 0 ) ( m ^ − m 0 ) ] + o P ( n − 1 / 2 ) . (100) \widehat{h}_\pi-\chi(\pi)=(P_n-P_0)\Gamma_0-P_0[(\widehat{w}-w_0)(\widehat{m}-m_0)]+o_{\mathbb{P}}(n^{-1/2}). \tag*{(100)} h π − χ ( π ) = ( P n − P 0 ) Γ 0 − P 0 [( w − w 0 ) ( m − m 0 )] + o P ( n − 1/2 ) . (100) For known weights, take q n = 0 q_n=0 q n = 0
Proof. Write v = φ S v=\varphi_S v = φ S D = H ^ − H 0 D=\widehat{H}-H_0 D = H − H 0
v π ( X ) = ∫ v ( a , X ) π ( d a ∣ X ) , b π = v π ( X ) − w 0 ( S ) v ( S ) , r 0 = φ Z ( Z ) − H 0 v ( S ) . v_\pi(X)=\int v(a,X)\pi(da\mid X),\qquad b_\pi=v_\pi(X)-w_0(S)v(S),\qquad r_0=\varphi_Z(Z)-H_0v(S). v π ( X ) = ∫ v ( a , X ) π ( d a ∣ X ) , b π = v π ( X ) − w 0 ( S ) v ( S ) , r 0 = φ Z ( Z ) − H 0 v ( S ) . Then P 0 b π = 0 P_0b_\pi=0 P 0 b π = 0 B π 2 : = P 0 ∥ b π ∥ 2 < ∞ B_\pi^2:=P_0\|b_\pi\|^2<\infty B π 2 := P 0 ∥ b π ∥ 2 < ∞ ∥ r 0 ∥ ≤ 2 κ Z \|r_0\|\le2\kappa_Z ∥ r 0 ∥ ≤ 2 κ Z
G n ( Γ η ^ − Γ 0 ) = D G n b π + G n [ q n r 0 ] − D G n [ q n v ] . (101) G_n(\Gamma_{\widehat{\eta}}-\Gamma_0)=DG_nb_\pi+G_n[q_nr_0]-DG_n[q_nv]. \tag*{(101)} G n ( Γ η − Γ 0 ) = D G n b π + G n [ q n r 0 ] − D G n [ q n v ] . (101) For either bounded Hilbert-valued function c = r 0 c=r_0 c = r 0 c = v c=v c = v U n , c = G n [ c ⊗ t ] U_{n,c}=G_n[c\otimes t] U n , c = G n [ c ⊗ t ] G n [ q c ] = U n , c q G_n[qc]=U_{n,c}q G n [ q c ] = U n , c q
E ∥ G n b π ∥ 2 = B π 2 , E ∥ U n , c ∥ H S 2 ≤ ∥ c ∥ ∞ 2 P 0 ∥ t ∥ 2 . \mathbb{E}\|G_nb_\pi\|^2=B_\pi^2,\qquad\mathbb{E}\|U_{n,c}\|_{\mathrm{HS}}^2\le\|c\|_\infty^2P_0\|t\|^2. E ∥ G n b π ∥ 2 = B π 2 , E ∥ U n , c ∥ HS 2 ≤ ∥ c ∥ ∞ 2 P 0 ∥ t ∥ 2 . Thus equation 97 proves the O P ( 1 ) O_{\mathbb{P}}(1) O P ( 1 )
For the stronger conclusion, fix R < ∞ R<\infty R < ∞ A R ( δ ) \mathcal{A}_R(\delta) A R ( δ ) ( D , q ) (D,q) ( D , q )
∥ D ∥ o p , ∥ q ∥ K w ≤ R , ∥ D v ∥ L 2 ( P S , 0 ; H Z ) , ∥ q ∥ L 2 ( P S , 0 ) ≤ δ . \|D\|_{\mathrm{op}},\ \|q\|_{\mathcal{K}_w}\le R,\qquad\|Dv\|_{L_2(P_{S,0};\mathcal{H}_Z)},\ \|q\|_{L_2(P_{S,0})}\le\delta. ∥ D ∥ op , ∥ q ∥ K w ≤ R , ∥ D v ∥ L 2 ( P S , 0 ; H Z ) , ∥ q ∥ L 2 ( P S , 0 ) ≤ δ . Let ( σ j , u j ) (\sigma_j,u_j) ( σ j , u j ) C S C_S C S Π J \Pi_J Π J J J J a J 2 = P 0 ∥ ( I − Π J ) b π ∥ 2 a_J^2=P_0\|(I-\Pi_J)b_\pi\|^2 a J 2 = P 0 ∥ ( I − Π J ) b π ∥ 2 v ( S ) v(S) v ( S ) v π ( X ) = E [ w 0 v ∣ X ] v_\pi(X)=\mathbb{E}[w_0v\mid X] v π ( X ) = E [ w 0 v ∣ X ] b π b_\pi b π a J → 0 a_J\to0 a J → 0 A R ( δ ) \mathcal{A}_R(\delta) A R ( δ ) ∥ D Π J ∥ o p ≤ δ / σ J \|D\Pi_J\|_{\mathrm{op}}\le\delta/\sqrt{\sigma_J} ∥ D Π J ∥ op ≤ δ / σ J
E sup ( D , q ) ∈ A R ( δ ) ∥ D G n b π ∥ ≤ δ B π σ J + R a J . \mathbb{E}\sup_{(D,q)\in\mathcal{A}_R(\delta)}\|DG_nb_\pi\|\le\frac{\delta B_\pi}{\sqrt{\sigma_J}}+Ra_J. E ( D , q ) ∈ A R ( δ ) sup ∥ D G n b π ∥ ≤ σ J δ B π + R a J . Similarly, let ( ρ j , e j ) (\rho_j,e_j) ( ρ j , e j ) T w = P 0 [ t ⊗ t ] T_w=P_0[t\otimes t] T w = P 0 [ t ⊗ t ] K K K Λ K \Lambda_K Λ K ∥ Λ K q ∥ ≤ δ / ρ K \|\Lambda_Kq\|\le\delta/\sqrt{\rho_K} ∥ Λ K q ∥ ≤ δ / ρ K
E ∥ U n , c ( I − Λ K ) ∥ H S 2 ≤ ∥ c ∥ ∞ 2 ∑ j > K ρ j , \mathbb{E}\|U_{n,c}(I-\Lambda_K)\|_{\mathrm{HS}}^2\le\|c\|_\infty^2\sum_{j>K}\rho_j, E ∥ U n , c ( I − Λ K ) ∥ HS 2 ≤ ∥ c ∥ ∞ 2 j > K ∑ ρ j , E ∗ sup ( D , q ) ∈ A R ( δ ) ∥ G n [ q c ] ∥ ≤ ∥ c ∥ ∞ { δ Tr ( T w ) ρ K + R ∑ j > K ρ j } . \mathbb{E}^{*}\sup_{(D,q)\in A_R(\delta)}\|\mathbb{G}_n[qc]\|\le\|c\|_\infty\left\{\delta\sqrt{\frac{\operatorname{Tr}(T_w)}{\rho_K}}+R\sqrt{\sum_{j>K}\rho_j}\right\}. E ∗ ( D , q ) ∈ A R ( δ ) sup ∥ G n [ q c ] ∥ ≤ ∥ c ∥ ∞ ⎩ ⎨ ⎧ δ ρ K Tr ( T w ) + R j > K ∑ ρ j ⎭ ⎬ ⎫ . All bounds are uniform in n n n J , K J,K J , K δ ↓ 0 \delta\downarrow0 δ ↓ 0
lim δ ↓ 0 lim sup n → ∞ P ∗ ( sup ( D , q ) ∈ A R ( δ ) ∥ G n ( Γ m 0 + D v , w 0 + q − Γ 0 ) ∥ > ϵ ) = 0 ( ϵ > 0 ) . \lim_{\delta\downarrow0}\limsup_{n\to\infty}\mathbb{P}^{*}\left(\sup_{(D,q)\in A_R(\delta)}\|\mathbb{G}_n(\Gamma_{m_0+Dv,w_0+q}-\Gamma_0)\|>\epsilon\right)=0\quad(\epsilon>0). δ ↓ 0 lim n → ∞ lim sup P ∗ ( ( D , q ) ∈ A R ( δ ) sup ∥ G n ( Γ m 0 + D v , w 0 + q − Γ 0 ) ∥ > ϵ ) = 0 ( ϵ > 0 ) . Norm tightness and L 2 L_2 L 2 o P ( 1 ) o_{\mathbb{P}}(1) o P ( 1 )
Lemma D.20 (Smooth parametric weights). *Retain the identification, kernel, fixed-policy, weight-moment, and outcome-operator conditions of Proposition D.19, and suppose ∥ m ^ − m 0 ∥ L 2 ( P S , 0 ; H Z ) = o P ( 1 ) \|\widehat m-m_0\|_{L_2(P_{S,0};\mathcal{H}_Z)}=o_{\mathbb{P}}(1) ∥ m − m 0 ∥ L 2 ( P S , 0 ; H Z ) = o P ( 1 ) w ^ = w θ ^ \widehat w=w_{\widehat\theta} w = w θ w θ 0 = w 0 w_{\theta_0}=w_0 w θ 0 = w 0 θ ∈ R p \theta\in\mathbb{R}^p θ ∈ R p p p p ∥ θ ^ − θ 0 ∥ 2 = O P ( n − 1 / 2 ) \|\widehat\theta-\theta_0\|_2=O_{\mathbb{P}}(n^{-1/2}) ∥ θ − θ 0 ∥ 2 = O P ( n − 1/2 ) U U U θ 0 \theta_0 θ 0 θ ↦ w θ ( s ) \theta\mapsto w_\theta(s) θ ↦ w θ ( s ) P S , 0 P_{S,0} P S , 0 s s s
∥ w ˙ θ 0 ∥ 2 ∈ L 2 ( P S , 0 ) , B ( s ) : = sup θ ∈ U ∥ w ¨ θ ( s ) ∥ o p ∈ L 2 ( P S , 0 ) . \|\dot w_{\theta_0}\|_2\in L_2(P_{S,0}),\qquad B(s):=\sup_{\theta\in U}\|\ddot w_\theta(s)\|_{\mathrm{op}}\in L_2(P_{S,0}). ∥ w ˙ θ 0 ∥ 2 ∈ L 2 ( P S , 0 ) , B ( s ) := θ ∈ U sup ∥ w ¨ θ ( s ) ∥ op ∈ L 2 ( P S , 0 ) . Then, even when both nuisances use all observations,
∥ w ^ − w 0 ∥ L 2 ( P S , 0 ) = O P ( n − 1 / 2 ) , ∥ G n ( Γ η ^ − Γ 0 ) ∥ H Z = o P ( 1 ) , \|\widehat w-w_0\|_{L_2(P_{S,0})}=O_{\mathbb{P}}(n^{-1/2}),\qquad\|\mathbb{G}_n(\Gamma_{\widehat\eta}-\Gamma_0)\|_{\mathcal{H}_Z}=o_{\mathbb{P}}(1), ∥ w − w 0 ∥ L 2 ( P S , 0 ) = O P ( n − 1/2 ) , ∥ G n ( Γ η − Γ 0 ) ∥ H Z = o P ( 1 ) , and
h ^ π − χ ( π ) = ( P n − P 0 ) Γ 0 + o P ( n − 1 / 2 ) . \widehat h_\pi-\chi(\pi)=(P_n-P_0)\Gamma_0+o_{\mathbb{P}}(n^{-1/2}). h π − χ ( π ) = ( P n − P 0 ) Γ 0 + o P ( n − 1/2 ) . Proof. Put Δ n = θ ^ − θ 0 \Delta_n=\widehat\theta-\theta_0 Δ n = θ − θ 0 θ ^ ∈ U \widehat\theta\in U θ ∈ U
q n = w ˙ θ 0 T Δ n + R n , ∣ R n ∣ ≤ 1 2 B ∥ Δ n ∥ 2 2 . q_n=\dot w_{\theta_0}^{\mathsf{T}}\Delta_n+R_n,\qquad|R_n|\le\tfrac{1}{2}B\|\Delta_n\|_2^2. q n = w ˙ θ 0 T Δ n + R n , ∣ R n ∣ ≤ 2 1 B ∥ Δ n ∥ 2 2 . The L 2 L_2 L 2 c c c
∥ G n [ q n c ] ∥ ≤ ∥ Δ n ∥ 2 ∑ j = 1 p ∥ G n [ w ˙ θ 0 , j c ] ∥ + 1 2 n ∥ Δ n ∥ 2 2 ∥ c ∥ ∞ ( P n B + P 0 B ) = O P ( n − 1 / 2 ) . \|\mathbb{G}_n[q_nc]\|\le\|\Delta_n\|_2\sum_{j=1}^{p}\|\mathbb{G}_n[\dot w_{\theta_0,j}c]\|+\tfrac{1}{2}\sqrt{n}\|\Delta_n\|_2^2\|c\|_\infty(P_nB+P_0B)=O_{\mathbb{P}}(n^{-1/2}). ∥ G n [ q n c ] ∥ ≤ ∥ Δ n ∥ 2 j = 1 ∑ p ∥ G n [ w ˙ θ 0 , j c ] ∥ + 2 1 n ∥ Δ n ∥ 2 2 ∥ c ∥ ∞ ( P n B + P 0 B ) = O P ( n − 1/2 ) . Here the fixed empirical processes have bounded second moments and P n B + P 0 B = O P ( 1 ) P_nB+P_0B=O_{\mathbb{P}}(1) P n B + P 0 B = O P ( 1 ) θ ^ \widehat\theta θ c = r 0 c=r_0 c = r 0 c = φ S c=\varphi_S c = φ S ∥ H ^ − H 0 ∥ o p = O P ( 1 ) \|\widehat H-H_0\|_{\mathrm{op}}=O_{\mathbb{P}}(1) ∥ H − H 0 ∥ op = O P ( 1 ) o P ( n − 1 / 2 ) o_{\mathbb{P}}(n^{-1/2}) o P ( n − 1/2 )
A concrete sufficient model is Gaussian logging regression A ∣ X ∼ N ( B 0 z ( X ) , Σ ) A\mid X\sim\mathcal{N}(B_0z(X),\Sigma) A ∣ X ∼ N ( B 0 z ( X ) , Σ ) z ( X ) z(X) z ( X ) E [ z ( X ) z ( X ) T ] \mathbb{E}[z(X)z(X)^{\mathsf{T}}] E [ z ( X ) z ( X ) T ] Σ ≻ 0 \Sigma\succ0 Σ ≻ 0 π ( ⋅ ∣ X ) = N ( μ π ( X ) , Σ ) \pi(\cdot\mid X)=\mathcal{N}(\mu_\pi(X),\Sigma) π ( ⋅ ∣ X ) = N ( μ π ( X ) , Σ ) μ π \mu_\pi μ π B ^ \widehat B B w B = π / p B w_B=\pi/p_B w B = π / p B n n n B 0 B_0 B 0 ∥ A ∥ \|A\| ∥ A ∥ exp ( C ∥ A ∥ + C ) \exp(C\|A\|+C) exp ( C ∥ A ∥ + C ) C C C P 0 w 0 2 < ∞ P_0w_0^2<\infty P 0 w 0 2 < ∞
Lemma D.18 and equation (97) verify the outcome conditions for the projected estimator, including with landmarks selected from the same data. For the weights, Proposition D.19 gives RKHS sufficient conditions, and Lemma D.20 gives a smooth parametric alternative. Also, ∥ Γ 0 ∥ ≤ κ Z ( 1 + 2 w 0 ) \|\Gamma_0\| \le\kappa_Z(1+2w_0) ∥ Γ 0 ∥ ≤ κ Z ( 1 + 2 w 0 ) (100) additionally requires the product of nuisance errors to be o P ( n − 1 / 2 ) o_{\mathrm{P}}(n^{-1/2}) o P ( n − 1/2 ) O P ( 1 ) O_{\mathrm{P}}(1) O P ( 1 )
Corollary D.21 (Projected Nyström DR embedding rate). Suppose Corollary D.17 holds and
∥ w ^ − w 0 ∥ L 2 ( P S , 0 ) = O P ( r w , n ) . \|\widehat{w}-w_0\|_{L_2(P_{\mathrm{S},0})}=O_{\mathrm{P}}(r_{w,n}). ∥ w − w 0 ∥ L 2 ( P S , 0 ) = O P ( r w , n ) . Assume further that
E ∥ Γ m 0 , w 0 ( O ) ∥ H Z 2 < ∞ \mathbb{E}\|\Gamma_{m_0,w_0}(O)\|_{\mathcal{H}_Z}^2 < \infty E ∥ Γ m 0 , w 0 ( O ) ∥ H Z 2 < ∞ and that
∥ ( P n − P 0 ) { Γ m ^ L n , λ n , w ^ − Γ m 0 , w 0 } ∥ H Z = O P ( n − 1 / 2 ) . \left\|(P_n-P_0)\left\{\Gamma_{\widehat{m}_{L_n,\lambda_n},\widehat{w}}-\Gamma_{m_0,w_0}\right\}\right\|_{\mathcal{H}_Z}=O_{\mathrm{P}}(n^{-1/2}). ( P n − P 0 ) { Γ m L n , λ n , w − Γ m 0 , w 0 } H Z = O P ( n − 1/2 ) . These two score conditions hold under the fixed-policy and ratio conditions of Proposition D.19 or Lemma D.20 , together with Lemma D.18 ; known weights require only P 0 w 0 2 < ∞ P_0w_0^2<\infty P 0 w 0 2 < ∞
∥ χ ^ D R , N y s ( π ) − χ ( π ) ∥ H Z = O P ( n − 1 / 2 + r w , n n − 1 2 ( 1 + γ ) ) . (102) \|\widehat{\chi}_{\mathrm{DR,Nys}}(\pi)-\chi(\pi)\|_{\mathcal{H}_Z}
=O_{\mathrm{P}}\left(n^{-1/2}+r_{w,n}n^{-\frac{1}{2(1+\gamma)}}\right). \tag*{(102)} ∥ χ DR , Nys ( π ) − χ ( π ) ∥ H Z = O P ( n − 1/2 + r w , n n − 2 ( 1 + γ ) 1 ) . (102) Proof. Apply equation (99) , the Hilbert-valued Cauchy–Schwarz bound for the product remainder, and Corollary D.17 .
Finite target-policy Monte Carlo approximation. The preceding results concern the projected Nyström doubly robust estimator defined using the exact target-policy integral. We now control the additional error induced by approximating this integral with J ≥ 1 J \ge1 J ≥ 1 [13 ] .
Write
m ^ n : = m ^ L n , λ n , χ ^ n : = χ ^ D R , N y s ( π ) . \widehat{m}_n := \widehat{m}_{L_n,\lambda_n}, \qquad\widehat{\chi}_n := \widehat{\chi}_{\mathrm{DR,Nys}}(\pi). m n := m L n , λ n , χ n := χ DR , Nys ( π ) . Here, χ ^ n \widehat{\chi}_n χ n i = 1 , … , n i=1,\ldots,n i = 1 , … , n
m ‾ n ( X i ) : = ∫ A m ^ n ( a , X i ) π ( d a ∣ X i ) , (103) \overline{m}_n(X_i):=\int_{\mathcal{A}}\widehat{m}_n(a,X_i)\,\pi(da\mid X_i), \tag*{(103)} m n ( X i ) := ∫ A m n ( a , X i ) π ( d a ∣ X i ) , (103) where the integral is a Bochner integral in H Z \mathcal{H}_Z H Z
χ ^ n = 1 n ∑ i = 1 n [ w ^ ( A i , X i ) { φ Z ( Z i ) − m ^ n ( A i , X i ) } + m ‾ n ( X i ) ] . (104) \widehat{\chi}_n=\frac{1}{n}\sum_{i=1}^{n}\left[\widehat{w}(A_i,X_i)\{\varphi_Z(Z_i)-\widehat{m}_n(A_i,X_i)\}+\overline{m}_n(X_i)\right]. \tag*{(104)} χ n = n 1 i = 1 ∑ n [ w ( A i , X i ) { φ Z ( Z i ) − m n ( A i , X i )} + m n ( X i ) ] . (104) Let F n \mathcal{F}_n F n m ^ n \widehat{m}_n m n w ^ \widehat{w} w χ ^ n \widehat{\chi}_n χ n F n \mathcal{F}_n F n
Conditional on F n \mathcal{F}_n F n
A i u π ∣ F n ∼ π ( ⋅ ∣ X i ) , i = 1 , … , n , u = 1 , … , J , (105) A_{iu}^{\pi}\mid\mathcal{F}_n\sim\pi(\cdot\mid X_i),\qquad i=1,\ldots,n,\quad u=1,\ldots,J, \tag*{(105)} A i u π ∣ F n ∼ π ( ⋅ ∣ X i ) , i = 1 , … , n , u = 1 , … , J , (105) independently across all pairs ( i , u ) (i,u) ( i , u )
m ‾ n , J ( X i ) : = 1 J ∑ u = 1 J m ^ n ( A i u π , X i ) , (106) \overline{m}_{n,J}(X_i):=\frac{1}{J}\sum_{u=1}^{J}\widehat{m}_n(A_{iu}^{\pi},X_i), \tag*{(106)} m n , J ( X i ) := J 1 u = 1 ∑ J m n ( A i u π , X i ) , (106) and the corresponding finite-J J J
χ ~ n , J : = 1 n ∑ i = 1 n [ w ^ ( A i , X i ) { φ Z ( Z i ) − m ^ n ( A i , X i ) } + m ‾ n , J ( X i ) ] . (107) \widetilde{\chi}_{n,J}:=\frac{1}{n}\sum_{i=1}^{n}\left[\widehat{w}(A_i,X_i)\{\varphi_Z(Z_i)-\widehat{m}_n(A_i,X_i)\}+\overline{m}_{n,J}(X_i)\right]. \tag*{(107)} χ n , J := n 1 i = 1 ∑ n [ w ( A i , X i ) { φ Z ( Z i ) − m n ( A i , X i )} + m n , J ( X i ) ] . (107) The exact and finite-J J J
χ ~ n , J − χ ^ n = 1 n ∑ i = 1 n { m ‾ n , J ( X i ) − m ‾ n ( X i ) } . (108) \widetilde{\chi}_{n,J}-\widehat{\chi}_n=\frac{1}{n}\sum_{i=1}^{n}\left\{\overline{m}_{n,J}(X_i)-\overline{m}_n(X_i)\right\}. \tag*{(108)} χ n , J − χ n = n 1 i = 1 ∑ n { m n , J ( X i ) − m n ( X i ) } . (108) Using the same target-policy draws and the same projected regression subspace, χ ~ n , J \widetilde{\chi}_{n,J} χ n , J
χ ~ D R , N y s , J ( π ) \widetilde{\chi}_{\mathrm{DR,Nys},J}(\pi) χ DR , Nys , J ( π ) in the matrix notation of Appendix § 12 .§ 2 .
Lemma D.22 (Finite-J J J Assume that, almost surely, for every i = 1 , … , n i=1,\ldots,n i = 1 , … , n
a ⟼ m ^ n ( a , X i ) a\longmapsto\widehat{m}_n(a,X_i) a ⟼ m n ( a , X i ) is strongly measurable and satisfies
∫ A ∥ m ^ n ( a , X i ) ∥ H Z 2 π ( d a ∣ X i ) < ∞ . (109) \int_{\mathcal{A}}\left\|\widehat{m}_n(a,X_i)\right\|_{\mathcal{H}_Z}^{2}\,\pi(da\mid X_i)<\infty. \tag*{(109)} ∫ A ∥ m n ( a , X i ) ∥ H Z 2 π ( d a ∣ X i ) < ∞. (109) Define the conditional average policy variance
V n : = 1 n ∑ i = 1 n ∫ A ∥ m ^ n ( a , X i ) − m ‾ n ( X i ) ∥ H Z 2 π ( d a ∣ X i ) . (110) V_n:=\frac{1}{n}\sum_{i=1}^{n}\int_{\mathcal{A}}\left\|\widehat{m}_n(a,X_i)-\overline{m}_n(X_i)\right\|_{\mathcal{H}_Z}^{2}\,\pi(da\mid X_i). \tag*{(110)} V n := n 1 i = 1 ∑ n ∫ A ∥ m n ( a , X i ) − m n ( X i ) ∥ H Z 2 π ( d a ∣ X i ) . (110) Then, almost surely,
E [ χ ~ n , J ∣ F n ] = χ ^ n , (111) \mathbb{E}\left[\widetilde{\chi}_{n,J}\mid\mathcal{F}_n\right]=\widehat{\chi}_n, \tag*{(111)} E [ χ n , J ∣ F n ] = χ n , (111) and
E [ ∥ χ ~ n , J − χ ^ n ∥ H Z 2 | F n ] = V n n J . (112) \mathbb{E}\left[\left\|\widetilde{\chi}_{n,J}-\widehat{\chi}_n\right\|_{\mathcal{H}_Z}^{2}\,\middle|\,\mathcal{F}_n\right]=\frac{V_n}{nJ}. \tag*{(112)} E [ ∥ χ n , J − χ n ∥ H Z 2 F n ] = n J V n . (112) Consequently, for every δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 )
Pr ( ∥ χ ~ n , J − χ ^ n ∥ H Z > V n n J δ | F n ) ≤ δ . (113) \Pr\left(\left\|\widetilde{\chi}_{n,J}-\widehat{\chi}_n\right\|_{\mathcal{H}_Z}>\sqrt{\frac{V_n}{nJ\delta}}\ \middle|\ \mathcal{F}_n\right)\leq\delta. \tag*{(113)} Pr ( ∥ χ n , J − χ n ∥ H Z > n J δ V n F n ) ≤ δ . (113) Proof. For i = 1 , … , n i=1,\ldots,n i = 1 , … , n u = 1 , … , J u=1,\ldots,J u = 1 , … , J
ξ i u : = m ^ n ( A i u π , X i ) − m ‾ n ( X i ) . \xi_{iu}:=\widehat{m}_n(A_{iu}^{\pi},X_i)-\overline{m}_n(X_i). ξ i u := m n ( A i u π , X i ) − m n ( X i ) . By the definition of the exact policy average, E [ ξ i u ∣ F n ] = 0 \mathbb{E}[\xi_{iu}\mid\mathcal{F}_n]=0 E [ ξ i u ∣ F n ] = 0 F n \mathcal{F}_n F n { ξ i u : i = 1 , … , n , u = 1 , … , J } \{\xi_{iu}:i=1,\ldots,n,\ u=1,\ldots,J\} { ξ i u : i = 1 , … , n , u = 1 , … , J }
Using equation (108) , we have
χ ~ n , J − χ ^ n = 1 n J ∑ i = 1 n ∑ u = 1 J ξ i u . (114) \widetilde{\chi}_{n,J}-\widehat{\chi}_n=\frac{1}{nJ}\sum_{i=1}^{n}\sum_{u=1}^{J}\xi_{iu}. \tag*{(114)} χ n , J − χ n = n J 1 i = 1 ∑ n u = 1 ∑ J ξ i u . (114) Conditional centering immediately gives equation (111) .
For two distinct index pairs ( i , u ) ≠ ( i ′ , u ′ ) (i,u)\ne(i',u') ( i , u ) = ( i ′ , u ′ )
E [ ⟨ ξ i u , ξ i ′ u ′ ⟩ H Z | F n ] = 0. \mathbb{E}\left[\left\langle\xi_{iu},\xi_{i'u'}\right\rangle_{\mathcal{H}_Z}\,\middle|\,\mathcal{F}_n\right]=0. E [ ⟨ ξ i u , ξ i ′ u ′ ⟩ H Z F n ] = 0. Expanding the squared norm in equation (114) therefore yields
E [ ∥ χ ~ n , J − χ ^ n ∥ H Z 2 | F n ] = 1 n 2 J 2 ∑ i = 1 n ∑ u = 1 J E [ ∥ ξ i u ∥ H Z 2 | F n ] = 1 n 2 J ∑ i = 1 n ∫ A ∥ m ^ n ( a , X i ) − m ‾ n ( X i ) ∥ H Z 2 π ( d a ∣ X i ) = V n n J . \begin{aligned}
\mathbb{E}\left[\left\|\widetilde{\chi}_{n,J}-\widehat{\chi}_n\right\|_{\mathcal{H}_Z}^{2}\,\middle|\,\mathcal{F}_n\right]
&=\frac{1}{n^{2}J^{2}}\sum_{i=1}^{n}\sum_{u=1}^{J}\mathbb{E}\left[\left\|\xi_{iu}\right\|_{\mathcal{H}_Z}^{2}\,\middle|\,\mathcal{F}_n\right]\\
&=\frac{1}{n^{2}J}\sum_{i=1}^{n}\int_{\mathcal{A}}\left\|\widehat{m}_n(a,X_i)-\overline{m}_n(X_i)\right\|_{\mathcal{H}_Z}^{2}\,\pi(da\mid X_i)\\
&=\frac{V_n}{nJ}.
\end{aligned} E [ ∥ χ n , J − χ n ∥ H Z 2 F n ] = n 2 J 2 1 i = 1 ∑ n u = 1 ∑ J E [ ∥ ξ i u ∥ H Z 2 F n ] = n 2 J 1 i = 1 ∑ n ∫ A ∥ m n ( a , X i ) − m n ( X i ) ∥ H Z 2 π ( d a ∣ X i ) = n J V n . This proves equation (112) . Finally, equation (113) follows by applying conditional Markov’s inequality to
∥ χ ~ n , J − χ ^ n ∥ H Z 2 . \left\|\widetilde{\chi}_{n,J}-\widehat{\chi}_n\right\|_{\mathcal{H}_Z}^{2}. ∥ χ n , J − χ n ∥ H Z 2 . Corollary D.23 (Finite-J J J Suppose the conditions of Corollary D.21 hold, and let ( J n ) n ≥ 1 (J_n)_{n\geq1} ( J n ) n ≥ 1
(i) Under the bounded-kernel and ridge assumptions already imposed,
∥ χ ~ n , J n − χ ^ n ∥ H Z = O P ( 1 n J n λ n ) . (115) \left\|\widetilde{\chi}_{n,J_n}-\widehat{\chi}_n\right\|_{\mathcal{H}_Z}
=O_{\mathrm{P}}\left(\frac{1}{\sqrt{nJ_n}\lambda_n}\right). \tag*{(115)} ∥ χ n , J n − χ n ∥ H Z = O P ( n J n λ n 1 ) . (115) Consequently,
∥ χ ~ n , J n − χ ( π ) ∥ H Z = O P ( n − 1 / 2 + r w , n n − 1 2 ( 1 + γ ) + 1 n J n λ n ) . (116) \left\|\widetilde{\chi}_{n,J_n}-\chi(\pi)\right\|_{\mathcal{H}_Z}
=O_{\mathrm{P}}\left(n^{-1/2}+r_{w,n}n^{-\frac{1}{2(1+\gamma)}}+\frac{1}{\sqrt{nJ_n}\lambda_n}\right). \tag*{(116)} ∥ χ n , J n − χ ( π ) ∥ H Z = O P ( n − 1/2 + r w , n n − 2 ( 1 + γ ) 1 + n J n λ n 1 ) . (116) For λ n = n − 1 / ( 1 + γ ) \lambda_n=n^{-1/(1+\gamma)} λ n = n − 1/ ( 1 + γ )
1 n J n λ n = J n − 1 / 2 n − γ 2 ( 1 + γ ) . (117) \frac{1}{\sqrt{nJ_n}\lambda_n}
=J_n^{-1/2}n^{-\frac{\gamma}{2(1+\gamma)}}. \tag*{(117)} n J n λ n 1 = J n − 1/2 n − 2 ( 1 + γ ) γ . (117) (ii) If, in addition,
V n = O P ( 1 ) , (118) V_n=O_{\mathrm{P}}(1), \tag*{(118)} V n = O P ( 1 ) , (118) then
∥ χ ~ n , J n − χ ^ n ∥ H Z = O P ( ( n J n ) − 1 / 2 ) . (119) \left\|\widetilde{\chi}_{n,J_n}-\widehat{\chi}_n\right\|_{\mathcal{H}_Z}
=O_{\mathrm{P}}\left((nJ_n)^{-1/2}\right). \tag*{(119)} ∥ χ n , J n − χ n ∥ H Z = O P ( ( n J n ) − 1/2 ) . (119) Consequently,
∥ χ ~ n , J n − χ ( π ) ∥ H Z = O P ( n − 1 / 2 + r w , n n − 1 2 ( 1 + γ ) + ( n J n ) − 1 / 2 ) . (120) \left\|\widetilde{\chi}_{n,J_n}-\chi(\pi)\right\|_{\mathcal{H}_Z}
=O_{\mathrm{P}}\left(n^{-1/2}+r_{w,n}n^{-\frac{1}{2(1+\gamma)}}+(nJ_n)^{-1/2}\right). \tag*{(120)} ∥ χ n , J n − χ ( π ) ∥ H Z = O P ( n − 1/2 + r w , n n − 2 ( 1 + γ ) 1 + ( n J n ) − 1/2 ) . (120) Since J n ≥ 1 J_n\geq1 J n ≥ 1
∥ χ ~ n , J n − χ ( π ) ∥ H Z = O P ( n − 1 / 2 + r w , n n − 1 2 ( 1 + γ ) ) . (121) \left\|\widetilde{\chi}_{n,J_n}-\chi(\pi)\right\|_{\mathcal{H}_Z}
=O_{\mathrm{P}}\left(n^{-1/2}+r_{w,n}n^{-\frac{1}{2(1+\gamma)}}\right). \tag*{(121)} ∥ χ n , J n − χ ( π ) ∥ H Z = O P ( n − 1/2 + r w , n n − 2 ( 1 + γ ) 1 ) . (121) Thus, under equation (118) , even a fixed number of target-policy draws per logged context does not change the exact-policy-integral statistical rate.
Proof. Because the zero operator is feasible in the projected ridge problem, optimality of H ^ L n , λ n \widehat{H}_{L_n,\lambda_n} H L n , λ n
λ n ∥ H ^ L n , λ n ∥ H S 2 ≤ 1 n ∑ i = 1 n ∥ φ Z ( Z i ) ∥ H Z 2 ≤ κ Z 2 . (122) \lambda_n\left\|\widehat{H}_{L_n,\lambda_n}\right\|_{\mathcal{H}_S}^{2}
\leq\frac{1}{n}\sum_{i=1}^{n}\left\|\varphi_Z(Z_i)\right\|_{\mathcal{H}_Z}^{2}
\leq\kappa_Z^{2}. \tag*{(122)} λ n H L n , λ n H S 2 ≤ n 1 i = 1 ∑ n ∥ φ Z ( Z i ) ∥ H Z 2 ≤ κ Z 2 . (122) Hence, for every s ∈ S s\in\mathcal{S} s ∈ S
∥ m ^ n ( s ) ∥ H Z = ∥ H ^ L n , λ n φ S ( s ) ∥ H Z ≤ ∥ H ^ L n , λ n ∥ H S ∥ φ S ( s ) ∥ H S ≤ κ S κ Z λ n . \left\|\widehat{m}_n(s)\right\|_{\mathcal{H}_Z}
=\left\|\widehat{H}_{L_n,\lambda_n}\varphi_S(s)\right\|_{\mathcal{H}_Z}
\leq\left\|\widehat{H}_{L_n,\lambda_n}\right\|_{\mathcal{H}_S}\left\|\varphi_S(s)\right\|_{\mathcal{H}_S}
\leq\frac{\kappa_S\kappa_Z}{\sqrt{\lambda_n}}. ∥ m n ( s ) ∥ H Z = H L n , λ n φ S ( s ) H Z ≤ H L n , λ n H S ∥ φ S ( s ) ∥ H S ≤ λ n κ S κ Z . Using the Hilbert-space variance identity,
∫ A ∥ m ^ n ( a , X i ) − m ‾ n ( X i ) ∥ H Z 2 π ( d a ∣ X i ) ≤ ∫ A ∥ m ^ n ( a , X i ) ∥ H Z 2 π ( d a ∣ X i ) , \int_{\mathcal{A}}\left\|\widehat{m}_n(a,X_i)-\overline{m}_n(X_i)\right\|_{\mathcal{H}_Z}^{2}\pi(da\mid X_i)
\leq\int_{\mathcal{A}}\left\|\widehat{m}_n(a,X_i)\right\|_{\mathcal{H}_Z}^{2}\pi(da\mid X_i), ∫ A ∥ m n ( a , X i ) − m n ( X i ) ∥ H Z 2 π ( d a ∣ X i ) ≤ ∫ A ∥ m n ( a , X i ) ∥ H Z 2 π ( d a ∣ X i ) , and therefore
V n ≤ κ S 2 κ Z 2 λ n . (123) V_n\leq\frac{\kappa_S^{2}\kappa_Z^{2}}{\lambda_n}. \tag*{(123)} V n ≤ λ n κ S 2 κ Z 2 . (123) Lemma D.22 now gives equation (115) . Combining it with Corollary D.21 and the triangle inequality gives equation (116) .
Under equation (118) , Lemma D.22 instead gives equation (119) . Combining this result with Corollary D.21 gives equation (120) . Finally,
( n J n ) − 1 / 2 ≤ n − 1 / 2 (nJ_n)^{-1/2}\leq n^{-1/2} ( n J n ) − 1/2 ≤ n − 1/2 because J n ≥ 1 J_n\geq1 J n ≥ 1 (121) .
Consistency of the generated particles
Throughout this section, write χ ( π ) = μ Q π \chi(\pi)=\mu_{Q_\pi} χ ( π ) = μ Q π h ^ π \widehat{h}_\pi h π h ~ π \widetilde{h}_\pi h π e n , J : = ∥ h ~ π − χ ( π ) ∥ H e_{n,J}:=\lVert\widetilde{h}_\pi-\chi(\pi)\rVert_{\mathcal{H}} e n , J := ∥ h π − χ ( π ) ∥ H Q ^ M ( r ) : = 1 M ∑ j = 1 M δ Z j ( r ) \widehat{Q}^{(r)}_M:=\frac{1}{M}\sum_{j=1}^{M}\delta_{Z_j^{(r)}} Q M ( r ) := M 1 ∑ j = 1 M δ Z j ( r )
Throughout this subsection, the particle bounds are understood conditionally on the target-norm event
E n : = { ∥ h ~ π ∥ H ≤ κ } . \mathcal{E}_n := \{\lVert\widetilde{h}_\pi\rVert_{\mathcal{H}}\le\kappa\}. E n := {∥ h π ∥ H ≤ κ } . For h ∈ H h\in\mathcal{H} h ∈ H
s h ( z ) : = ⟨ h , φ ( z ) ⟩ H , v h ( Q , z ) = ∇ s h ( z ) − ∫ ∇ 1 k ( z , z ′ ) Q ( d z ′ ) . s_h(z):=\langle h,\varphi(z)\rangle_{\mathcal{H}}, \qquad v_h(Q,z)=\nabla s_h(z)-\int\nabla_1 k(z,z')Q(\mathrm{d}z'). s h ( z ) := ⟨ h , φ ( z ) ⟩ H , v h ( Q , z ) = ∇ s h ( z ) − ∫ ∇ 1 k ( z , z ′ ) Q ( d z ′ ) . Let D \mathcal{D} D F r \mathcal{F}_r F r r r r
Z j ( r + 1 ) = Z j ( r ) + τ v h ~ π ( Q ^ M ( r ) , Z j ( r ) + β r U j ( r ) ) , j = 1 , … , M , Z_j^{(r+1)}=Z_j^{(r)}+\tau v_{\widetilde{h}_\pi}\left(\widehat{Q}^{(r)}_M,Z_j^{(r)}+\beta_rU_j^{(r)}\right), \qquad j=1,\ldots,M, Z j ( r + 1 ) = Z j ( r ) + τ v h π ( Q M ( r ) , Z j ( r ) + β r U j ( r ) ) , j = 1 , … , M , where, conditional on F r \mathcal{F}_r F r U j ( r ) U_j^{(r)} U j ( r ) R d Z \mathbb{R}^{d_Z} R d Z
Assumption D.25 (Smooth particle dynamics). The feature map φ \varphi φ H \mathcal{H} H κ ≥ 1 \kappa\ge1 κ ≥ 1 η r ≡ τ > 0 \eta_r\equiv\tau>0 η r ≡ τ > 0 β r = c β ( r + 1 ) − α \beta_r=c_\beta(r+1)^{-\alpha} β r = c β ( r + 1 ) − α
12 τ κ 2 d Z ≤ 1 , 0 < c β ≤ 1 2 d Z , 0 ≤ α < 1 2 . 12\tau\kappa^2d_Z\le1,\qquad0<c_\beta\le\frac{1}{2\sqrt{d_Z}},\qquad0\le\alpha<\frac{1}{2}. 12 τ κ 2 d Z ≤ 1 , 0 < c β ≤ 2 d Z 1 , 0 ≤ α < 2 1 . Assumption D.26 (Noisy gradient dominance). For r = 0 , … , R − 1 r=0,\ldots,R-1 r = 0 , … , R − 1 h ~ π \widetilde{h}_\pi h π
1 M ∑ j = 1 M E [ ∥ v χ ( π ) ( Q ^ M ( r ) , Z j ( r ) + β r U j ( r ) ) ∥ 2 2 | F r ] ≥ 4 κ 2 d Z β r 2 MMD k 2 ( Q ^ M ( r ) , Q π ) . \frac{1}{M}\sum_{j=1}^{M}\mathbb{E}\left[\left\lVert v_{\chi(\pi)}\left(\widehat{Q}^{(r)}_M,Z_j^{(r)}+\beta_rU_j^{(r)}\right)\right\rVert_2^2\middle|\mathcal{F}_r\right]\ge4\kappa^2d_Z\beta_r^2\operatorname{MMD}_k^2\left(\widehat{Q}^{(r)}_M,Q^\pi\right). M 1 j = 1 ∑ M E [ v χ ( π ) ( Q M ( r ) , Z j ( r ) + β r U j ( r ) ) 2 2 F r ] ≥ 4 κ 2 d Z β r 2 MMD k 2 ( Q M ( r ) , Q π ) . Lemma D.27 (Descent lemma in expectation). Under Assumptions D.25 and D.26,
E [ ∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H 2 | F r ] ≤ ( 1 − τ κ 2 d Z 8 β r 2 ) ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 + τ κ 2 ( 9 2 d Z β r 2 + 3 4 ) e n , J 2 . \mathbb{E}\left[\left\lVert\mu_{\widehat{Q}^{(r+1)}_M}-\widetilde{h}_\pi\right\rVert_{\mathcal{H}}^2\middle|\mathcal{F}_r\right]\le\left(1-\frac{\tau\kappa^2d_Z}{8}\beta_r^2\right)\left\lVert\mu_{\widehat{Q}^{(r)}_M}-\widetilde{h}_\pi\right\rVert_{\mathcal{H}}^2+\tau\kappa^2\left(\frac{9}{2}d_Z\beta_r^2+\frac{3}{4}\right)e_{n,J}^2. E [ μ Q M ( r + 1 ) − h π H 2 F r ] ≤ ( 1 − 8 τ κ 2 d Z β r 2 ) μ Q M ( r ) − h π H 2 + τ κ 2 ( 2 9 d Z β r 2 + 4 3 ) e n , J 2 . In particular, with a : = τ κ 2 d Z c β 2 8 a:=\frac{\tau\kappa^2d_Zc_\beta^2}{8} a := 8 τ κ 2 d Z c β 2 B : = τ κ 2 ( 9 2 d Z c β 2 + 3 4 ) B:=\tau\kappa^2\left(\frac{9}{2}d_Zc_\beta^2+\frac{3}{4}\right) B := τ κ 2 ( 2 9 d Z c β 2 + 4 3 ) a r : = a ( r + 1 ) − 2 α a_r:=a(r+1)^{-2\alpha} a r := a ( r + 1 ) − 2 α
( 1 − a r ) ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 + B e n , J 2 . (1-a_r)\left\lVert\mu_{\widehat{Q}^{(r)}_M}-\widetilde{h}_\pi\right\rVert_{\mathcal{H}}^2+Be_{n,J}^2. ( 1 − a r ) μ Q M ( r ) − h π H 2 + B e n , J 2 . Lemma D.28 (Concentration of particles after noise injection). Under Assumption D.25, for any q > 0 q>0 q > 0 F r \mathcal{F}_r F r 1 − e − q 1-e^{-q} 1 − e − q
∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H − E [ ∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H | F r ] ≤ 2 τ κ 3 β r 2 q M . \left\lVert\mu_{\widehat{Q}^{(r+1)}_M}-\widetilde{h}_\pi\right\rVert_{\mathcal{H}}-\mathbb{E}\left[\left\lVert\mu_{\widehat{Q}^{(r+1)}_M}-\widetilde{h}_\pi\right\rVert_{\mathcal{H}}\middle|\mathcal{F}_r\right]\le2\tau\kappa^3\beta_r\sqrt{\frac{2q}{M}}. μ Q M ( r + 1 ) − h π H − E [ μ Q M ( r + 1 ) − h π H F r ] ≤ 2 τ κ 3 β r M 2 q . Proof of Theorem 5.3. Condition on D \mathcal{D} D ∥ h ~ π ∥ H ≤ κ \lVert\widetilde{h}_{\pi}\rVert_{\mathcal{H}}\leq\kappa ∥ h π ∥ H ≤ κ ∥ μ Q ^ M ( r ) − h ~ π ∥ H ≤ 2 κ \lVert\mu_{\widehat{Q}^{(r)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}\leq2\kappa ∥ μ Q M ( r ) − h π ∥ H ≤ 2 κ a a a B B B a r a_r a r Lemma D.27 , so that 0 < a r < 1 0<a_r<1 0 < a r < 1
Fix δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 ) Lemma D.28 with q = log ( R / δ ) q=\log(R/\delta) q = log ( R / δ ) Lemma D.27 , and a conditional union bound, shows that, with conditional probability at least 1 − δ 1-\delta 1 − δ r = 0 , … , R − 1 r=0,\ldots,R-1 r = 0 , … , R − 1
∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H ≤ ( 1 − a r ) ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 + B e n , J 2 + 2 2 τ κ 3 c β ( r + 1 ) − α log ( R / δ ) M . \lVert\mu_{\widehat{Q}^{(r+1)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}\leq\sqrt{(1-a_r)\lVert\mu_{\widehat{Q}^{(r)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}^2+Be_{n,J}^2+2\sqrt{2}\,\tau\kappa^3c_{\beta}(r+1)^{-\alpha}\sqrt{\frac{\log(R/\delta)}{M}}}. ∥ μ Q M ( r + 1 ) − h π ∥ H ≤ ( 1 − a r ) ∥ μ Q M ( r ) − h π ∥ H 2 + B e n , J 2 + 2 2 τ κ 3 c β ( r + 1 ) − α M log ( R / δ ) . Work on this event. The map x ↦ ( 1 − a r ) x 2 + B e n , J 2 x\mapsto\sqrt{(1-a_r)x^2+Be_{n,J}^2} x ↦ ( 1 − a r ) x 2 + B e n , J 2 1 − a r \sqrt{1-a_r} 1 − a r [ 0 , ∞ ) [0,\infty) [ 0 , ∞ )
∥ μ Q ^ M ( R ) − h ~ π ∥ H ≤ [ ∥ μ Q ^ M ( 0 ) − h ~ π ∥ H 2 ∏ r = 0 R − 1 ( 1 − a r ) + B e n , J 2 ∑ r = 0 R − 1 ∏ ℓ = r + 1 R − 1 ( 1 − a ℓ ) ] 1 / 2 \lVert\mu_{\widehat{Q}^{(R)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}\leq\left[\lVert\mu_{\widehat{Q}^{(0)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}^2\prod_{r=0}^{R-1}(1-a_r)+Be_{n,J}^2\sum_{r=0}^{R-1}\prod_{\ell=r+1}^{R-1}(1-a_\ell)\right]^{1/2} ∥ μ Q M ( R ) − h π ∥ H ≤ [ ∥ μ Q M ( 0 ) − h π ∥ H 2 r = 0 ∏ R − 1 ( 1 − a r ) + B e n , J 2 r = 0 ∑ R − 1 ℓ = r + 1 ∏ R − 1 ( 1 − a ℓ ) ] 1/2 + 2 2 τ κ 3 c β log ( R / δ ) M ∑ r = 0 R − 1 ( r + 1 ) − α ∏ ℓ = r + 1 R − 1 1 − a ℓ . \qquad+2\sqrt{2}\,\tau\kappa^3c_{\beta}\sqrt{\frac{\log(R/\delta)}{M}}\sum_{r=0}^{R-1}(r+1)^{-\alpha}\prod_{\ell=r+1}^{R-1}\sqrt{1-a_\ell}. + 2 2 τ κ 3 c β M log ( R / δ ) r = 0 ∑ R − 1 ( r + 1 ) − α ℓ = r + 1 ∏ R − 1 1 − a ℓ . Since a ℓ ≥ a R − 2 α a_\ell\geq aR^{-2\alpha} a ℓ ≥ a R − 2 α ℓ < R \ell<R ℓ < R ∑ r = 0 R − 1 ∏ ℓ = r + 1 R − 1 ( 1 − a ℓ ) ≤ R 2 α / a \sum_{r=0}^{R-1}\prod_{\ell=r+1}^{R-1}(1-a_\ell)\leq R^{2\alpha}/a ∑ r = 0 R − 1 ∏ ℓ = r + 1 R − 1 ( 1 − a ℓ ) ≤ R 2 α / a 1 − x ≤ e − x / 2 \sqrt{1-x}\leq e^{-x/2} 1 − x ≤ e − x /2 ∑ r = 0 R − 1 ( r + 1 ) − 2 α ≥ ∫ 1 R + 1 x − 2 α d x \sum_{r=0}^{R-1}(r+1)^{-2\alpha}\geq\int_1^{R+1}x^{-2\alpha}\,\mathrm{d}x ∑ r = 0 R − 1 ( r + 1 ) − 2 α ≥ ∫ 1 R + 1 x − 2 α d x ∏ r = 0 R − 1 1 − a r ≤ exp [ − a { ( R + 1 ) 1 − 2 α − 1 } / { 2 ( 1 − 2 α ) } ] \prod_{r=0}^{R-1}\sqrt{1-a_r}\leq\exp\left[-a\{(R+1)^{1-2\alpha}-1\}/\{2(1-2\alpha)\}\right] ∏ r = 0 R − 1 1 − a r ≤ exp [ − a {( R + 1 ) 1 − 2 α − 1 } / { 2 ( 1 − 2 α )} ] ( r + 1 ) − α ≤ R α a r / a (r+1)^{-\alpha}\leq R^\alpha a_r/a ( r + 1 ) − α ≤ R α a r / a a r ≤ 2 ( 1 − 1 − a r ) a_r\leq2(1-\sqrt{1-a_r}) a r ≤ 2 ( 1 − 1 − a r ) ∑ r = 0 R − 1 ( r + 1 ) − α ∏ ℓ = r + 1 R − 1 1 − a ℓ ≤ 2 R α / a \sum_{r=0}^{R-1}(r+1)^{-\alpha}\prod_{\ell=r+1}^{R-1}\sqrt{1-a_\ell}\leq2R^\alpha/a ∑ r = 0 R − 1 ( r + 1 ) − α ∏ ℓ = r + 1 R − 1 1 − a ℓ ≤ 2 R α / a
Combining these estimates with x + y ≤ x + y \sqrt{x+y}\leq\sqrt{x}+\sqrt{y} x + y ≤ x + y ∥ μ Q ^ M ( 0 ) − h ~ π ∥ H ≤ 2 κ \lVert\mu_{\widehat{Q}^{(0)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}\leq2\kappa ∥ μ Q M ( 0 ) − h π ∥ H ≤ 2 κ MMD k ( Q ^ M ( R ) , Q π ) ≤ ∥ μ Q ^ M ( R ) − h ~ π ∥ H + e n , J \operatorname{MMD}_k(\widehat{Q}^{(R)}_{M},Q^\pi)\leq\lVert\mu_{\widehat{Q}^{(R)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}+e_{n,J} MMD k ( Q M ( R ) , Q π ) ≤ ∥ μ Q M ( R ) − h π ∥ H + e n , J
MMD k ( Q ^ M ( R ) , Q π ) ≤ C 0 e − c R 1 − 2 α + ( 1 + C 1 R α ) e n , J + C 2 R α log ( R / δ ) M . \operatorname{MMD}_k(\widehat{Q}^{(R)}_{M},Q^\pi)\leq C_0e^{-cR^{1-2\alpha}}+(1+C_1R^\alpha)e_{n,J}+C_2R^\alpha\sqrt{\frac{\log(R/\delta)}{M}}. MMD k ( Q M ( R ) , Q π ) ≤ C 0 e − c R 1 − 2 α + ( 1 + C 1 R α ) e n , J + C 2 R α M log ( R / δ ) . where one may take c = a / { 2 ( 1 − 2 α ) } c=a/\{2(1-2\alpha)\} c = a / { 2 ( 1 − 2 α )} C 0 = 2 κ e c C_0=2\kappa e^c C 0 = 2 κ e c C 1 = B / a C_1=\sqrt{B/a} C 1 = B / a C 2 = 4 2 τ κ 3 c β / a C_2=4\sqrt{2}\,\tau\kappa^3c_\beta/a C 2 = 4 2 τ κ 3 c β / a
Remark D.29 (Relation to the stationary MMD points result). The auxiliary RKHS Lipschitz argument and discrete recursion techniques from the stationary MMD points analysis [27 ] apply here. Its descent lemma cannot be applied unchanged: the particle update uses h ~ π \widetilde{h}_{\pi} h π Q π Q^\pi Q π e n , J 2 e_{n,J}^2 e n , J 2 Lemma D.27 . The concentration step is stated for the updated particles in Lemma D.28 . Solving the squared recursion before taking a square root gives a propagated target error of order R α e n , J R^\alpha e_{n,J} R α e n , J R α log ( R / δ ) / M R^\alpha\sqrt{\log(R/\delta)/M} R α log ( R / δ ) / M
Proof of Lemma D.27 . Fix r r r F r \mathcal{F}_r F r
∥ ∇ h ~ π ( Q ^ M ( r ) , z ) − ∇ h ~ π ( Q ^ M ( r ) , z ′ ) ∥ 2 ≤ κ ∥ μ Q ^ M ( r ) − h ~ π ∥ H ∥ z − z ′ ∥ 2 , \lVert\nabla\widetilde{h}_{\pi}(\widehat{Q}^{(r)}_{M},z)-\nabla\widetilde{h}_{\pi}(\widehat{Q}^{(r)}_{M},z')\rVert_2\leq\kappa\lVert\mu_{\widehat{Q}^{(r)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}\lVert z-z'\rVert_2, ∥ ∇ h π ( Q M ( r ) , z ) − ∇ h π ( Q M ( r ) , z ′ ) ∥ 2 ≤ κ ∥ μ Q M ( r ) − h π ∥ H ∥ z − z ′ ∥ 2 , ∥ ∇ h ~ π ( Q ^ M ( r ) , z ) − ∇ χ ( π ) ( Q ^ M ( r ) , z ) ∥ 2 ≤ κ e n , J . \lVert\nabla\widetilde{h}_{\pi}(\widehat{Q}^{(r)}_{M},z)-\nabla\chi_{(\pi)}(\widehat{Q}^{(r)}_{M},z)\rVert_2\leq\kappa e_{n,J}. ∥ ∇ h π ( Q M ( r ) , z ) − ∇ χ ( π ) ( Q M ( r ) , z ) ∥ 2 ≤ κ e n , J . These bounds depend on the target only through its RKHS residual and do not require h ~ π \widetilde{h}_{\pi} h π
The gradient of ∥ M − 1 ∑ j = 1 M φ ( z j ) − h ~ π ∥ H 2 \lVert M^{-1}\sum_{j=1}^{M}\varphi(z_j)-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}^2 ∥ M − 1 ∑ j = 1 M φ ( z j ) − h π ∥ H 2 j j j − 2 ∇ h ~ π ( Q ^ M ( r ) , Z j ( r ) ) / M -2\nabla\widetilde{h}_{\pi}(\widehat{Q}^{(r)}_{M},Z_j^{(r)})/M − 2∇ h π ( Q M ( r ) , Z j ( r ) ) / M 2 κ 2\kappa 2 κ κ \kappa κ u 1 : M u_{1:M} u 1 : M 6 κ 2 M − 1 ∑ j = 1 M ∥ u j ∥ 2 2 6\kappa^2M^{-1}\sum_{j=1}^{M}\lVert u_j\rVert_2^2 6 κ 2 M − 1 ∑ j = 1 M ∥ u j ∥ 2 2
∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H 2 ≤ ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 − 2 τ M ∑ j = 1 M ⟨ ∇ h ~ π ( Q ^ M ( r ) , Z j ( r ) ) , ∇ h ~ π ( Q ^ M ( r ) , Z j ( r ) ) + β r U j ( r ) ⟩ + 3 τ 2 κ 2 M ∑ j = 1 M ∥ v h ~ π ( Q ^ M ( r ) , Z j ( r ) + β r U j ( r ) ) ∥ 2 2 . \begin{aligned}
\lVert\mu_{\widehat{Q}^{(r+1)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}^2\leq\lVert\mu_{\widehat{Q}^{(r)}_{M}}-\widetilde{h}_{\pi}\rVert_{\mathcal{H}}^2-\frac{2\tau}{M}\sum_{j=1}^{M}\left\langle\nabla\widetilde{h}_{\pi}(\widehat{Q}^{(r)}_{M},Z_j^{(r)}),\nabla\widetilde{h}_{\pi}(\widehat{Q}^{(r)}_{M},Z_j^{(r)})+\beta_rU_j^{(r)}\right\rangle\\
&+\frac{3\tau^{2}\kappa^{2}}{M}\sum_{j=1}^{M}\left\|v_{\tilde h_{\pi}}\left(\widehat Q_{M}^{(r)},Z_{j}^{(r)}+\beta_{r}U_{j}^{(r)}\right)\right\|_{2}^{2}.
\end{aligned} ∥ μ Q M ( r + 1 ) − h π ∥ H 2 ≤ ∥ μ Q M ( r ) − h π ∥ H 2 − M 2 τ j = 1 ∑ M ⟨ ∇ h π ( Q M ( r ) , Z j ( r ) ) , ∇ h π ( Q M ( r ) , Z j ( r ) ) + β r U j ( r ) ⟩ + M 3 τ 2 κ 2 j = 1 ∑ M v h ~ π ( Q M ( r ) , Z j ( r ) + β r U j ( r ) ) 2 2 . By the first Lipschitz bound and E [ ∥ U j ( r ) ∥ 2 2 ∣ F r ] = d Z \mathbb{E}[\|U_{j}^{(r)}\|_{2}^{2}\mid\mathcal{F}_{r}]=d_{Z} E [ ∥ U j ( r ) ∥ 2 2 ∣ F r ] = d Z κ 2 d Z β r 2 ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 \kappa^{2}d_{Z}\beta_{r}^{2}\|\mu_{\widehat Q_{M}^{(r)}}-\widetilde h_{\pi}\|_{\mathcal{H}}^{2} κ 2 d Z β r 2 ∥ μ Q M ( r ) − h π ∥ H 2 2 ⟨ a , b ⟩ = ∥ a ∥ 2 2 + ∥ b ∥ 2 2 − ∥ a − b ∥ 2 2 2\langle a,b\rangle=\|a\|_{2}^{2}+\|b\|_{2}^{2}-\|a-b\|_{2}^{2} 2 ⟨ a , b ⟩ = ∥ a ∥ 2 2 + ∥ b ∥ 2 2 − ∥ a − b ∥ 2 2 3 τ κ 2 ≤ 1 / 4 3\tau\kappa^{2}\le1/4 3 τ κ 2 ≤ 1/4
E [ ∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H 2 ∣ F r ] ≤ ( 1 + τ κ 2 d Z β r 2 ) ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 − 3 τ 4 M ∑ j = 1 M E [ ∥ v h ~ π ( Q ^ M ( r ) , Z j ( r ) + β r U j ( r ) ) ∥ 2 2 ∣ F r ] . \mathbb{E}\left[\|\mu_{\widehat Q_{M}^{(r+1)}}-\widetilde h_{\pi}\|_{\mathcal{H}}^{2}\mid\mathcal{F}_{r}\right]\leq\left(1+\tau\kappa^{2}d_{Z}\beta_{r}^{2}\right)\|\mu_{\widehat Q_{M}^{(r)}}-\widetilde h_{\pi}\|_{\mathcal{H}}^{2}-\frac{3\tau}{4M}\sum_{j=1}^{M}\mathbb{E}\left[\left\|v_{\tilde h_{\pi}}\left(\widehat Q_{M}^{(r)},Z_{j}^{(r)}+\beta_{r}U_{j}^{(r)}\right)\right\|_{2}^{2}\mid\mathcal{F}_{r}\right]. E [ ∥ μ Q M ( r + 1 ) − h π ∥ H 2 ∣ F r ] ≤ ( 1 + τ κ 2 d Z β r 2 ) ∥ μ Q M ( r ) − h π ∥ H 2 − 4 M 3 τ j = 1 ∑ M E [ v h ~ π ( Q M ( r ) , Z j ( r ) + β r U j ( r ) ) 2 2 ∣ F r ] . The velocity perturbation bound, the inequality ∥ a + b ∥ 2 2 ≥ 1 2 ∥ a ∥ 2 2 − ∥ b ∥ 2 2 \|a+b\|_{2}^{2}\geq\frac{1}{2}\|a\|_{2}^{2}-\|b\|_{2}^{2} ∥ a + b ∥ 2 2 ≥ 2 1 ∥ a ∥ 2 2 − ∥ b ∥ 2 2 Assumption D.26 imply
1 M ∑ j = 1 M E [ ∥ v h ~ π ( Q ^ M ( r ) , Z j ( r ) + β r U j ( r ) ) ∥ 2 2 ∣ F r ] ≥ 2 κ 2 d Z β r 2 MMD k 2 ( Q ^ M ( r ) , Q π ) − κ 2 e n , J 2 ≥ 3 2 κ 2 d Z β r 2 ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 − ( 6 κ 2 d Z β r 2 + κ 2 ) e n , J 2 , \begin{aligned}
\frac{1}{M}\sum_{j=1}^{M}\mathbb{E}\left[\left\|v_{\tilde h_{\pi}}\left(\widehat Q_{M}^{(r)},Z_{j}^{(r)}+\beta_{r}U_{j}^{(r)}\right)\right\|_{2}^{2}\mid\mathcal{F}_{r}\right]\geq2\kappa^{2}d_{Z}\beta_{r}^{2}\operatorname{MMD}_{k}^{2}\left(\widehat Q_{M}^{(r)},Q^{\pi}\right)-\kappa^{2}e_{n,J}^{2} \\
&\geq\frac{3}{2}\kappa^{2}d_{Z}\beta_{r}^{2}\|\mu_{\widehat Q_{M}^{(r)}}-\widetilde h_{\pi}\|_{\mathcal{H}}^{2}-(6\kappa^{2}d_{Z}\beta_{r}^{2}+\kappa^{2})e_{n,J}^{2},
\end{aligned} M 1 j = 1 ∑ M E [ v h ~ π ( Q M ( r ) , Z j ( r ) + β r U j ( r ) ) 2 2 ∣ F r ] ≥ 2 κ 2 d Z β r 2 MMD k 2 ( Q M ( r ) , Q π ) − κ 2 e n , J 2 ≥ 2 3 κ 2 d Z β r 2 ∥ μ Q M ( r ) − h π ∥ H 2 − ( 6 κ 2 d Z β r 2 + κ 2 ) e n , J 2 , where the second inequality uses MMD k 2 ( Q ^ M ( r ) , Q π ) ≥ 3 4 ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 − 3 e n , J 2 \operatorname{MMD}_{k}^{2}(\widehat Q_{M}^{(r)},Q^{\pi})\geq\frac{3}{4}\|\mu_{\widehat Q_{M}^{(r)}}-\widetilde h_{\pi}\|_{\mathcal{H}}^{2}-3e_{n,J}^{2} MMD k 2 ( Q M ( r ) , Q π ) ≥ 4 3 ∥ μ Q M ( r ) − h π ∥ H 2 − 3 e n , J 2
E [ ∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H 2 ∣ F r ] ≤ ( 1 − τ κ 2 d Z 8 β r 2 ) ∥ μ Q ^ M ( r ) − h ~ π ∥ H 2 + τ κ 2 ( 9 2 d Z β r 2 + 3 4 ) e n , J 2 . \mathbb{E}\left[\|\mu_{\widehat Q_{M}^{(r+1)}}-\widetilde h_{\pi}\|_{\mathcal{H}}^{2}\mid\mathcal{F}_{r}\right]\leq\left(1-\frac{\tau\kappa^{2}d_{Z}}{8}\beta_{r}^{2}\right)\|\mu_{\widehat Q_{M}^{(r)}}-\widetilde h_{\pi}\|_{\mathcal{H}}^{2}+\tau\kappa^{2}\left(\frac{9}{2}d_{Z}\beta_{r}^{2}+\frac{3}{4}\right)e_{n,J}^{2}. E [ ∥ μ Q M ( r + 1 ) − h π ∥ H 2 ∣ F r ] ≤ ( 1 − 8 τ κ 2 d Z β r 2 ) ∥ μ Q M ( r ) − h π ∥ H 2 + τ κ 2 ( 2 9 d Z β r 2 + 4 3 ) e n , J 2 . Finally, β r 2 = c β 2 ( r + 1 ) − 2 α ≤ c β 2 \beta_{r}^{2}=c_{\beta}^{2}(r+1)^{-2\alpha}\leq c_{\beta}^{2} β r 2 = c β 2 ( r + 1 ) − 2 α ≤ c β 2 a r a_{r} a r B B B
Proof of Lemma D.28 . Condition on F r \mathcal{F}_{r} F r v h ~ π ( Q ^ M ( r ) , ⋅ ) v_{\tilde h_{\pi}}(\widehat Q_{M}^{(r)},\cdot) v h ~ π ( Q M ( r ) , ⋅ ) κ ∥ μ Q ^ M ( r ) − h ~ π ∥ H \kappa\|\mu_{\widehat Q_{M}^{(r)}}-\widetilde h_{\pi}\|_{\mathcal{H}} κ ∥ μ Q M ( r ) − h π ∥ H φ \varphi φ κ \kappa κ U 1 : M ( r ) U_{1:M}^{(r)} U 1 : M ( r ) ∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H \|\mu_{\widehat Q_{M}^{(r+1)}}-\widetilde h_{\pi}\|_{\mathcal{H}} ∥ μ Q M ( r + 1 ) − h π ∥ H τ κ 2 β r ∥ μ Q ^ M ( r ) − h ~ π ∥ H / M ≤ 2 τ κ 3 β r / M \tau\kappa^{2}\beta_{r}\|\mu_{\widehat Q_{M}^{(r)}}-\widetilde h_{\pi}\|_{\mathcal{H}}/\sqrt{M}\leq2\tau\kappa^{3}\beta_{r}/\sqrt{M} τ κ 2 β r ∥ μ Q M ( r ) − h π ∥ H / M ≤ 2 τ κ 3 β r / M q > 0 q>0 q > 0 1 − e − q 1-e^{-q} 1 − e − q
∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H − E [ ∥ μ Q ^ M ( r + 1 ) − h ~ π ∥ H ∣ F r ] ≤ 2 τ κ 3 β r 2 q M . \|\mu_{\widehat Q_{M}^{(r+1)}}-\widetilde h_{\pi}\|_{\mathcal{H}}-\mathbb{E}\left[\|\mu_{\widehat Q_{M}^{(r+1)}}-\widetilde h_{\pi}\|_{\mathcal{H}}\mid\mathcal{F}_{r}\right]\leq2\tau\kappa^{3}\beta_{r}\sqrt{\frac{2q}{M}}. ∥ μ Q M ( r + 1 ) − h π ∥ H − E [ ∥ μ Q M ( r + 1 ) − h π ∥ H ∣ F r ] ≤ 2 τ κ 3 β r M 2 q . This proves the claim.
Numerical-target rate and consistency
Proof of Proposition 5.5 . Apply the exact expansion equation (99) to the exact-integral estimator. Its oracle term is O p ( n − 1 / 2 ) O_{p}(n^{-1/2}) O p ( n − 1/2 ) Assumption 5.4 , and its product remainder is O p ( r m , n r w , n ) O_{p}(r_{m,n}r_{w,n}) O p ( r m , n r w , n ) V n / ( n J ) V_{n}/(nJ) V n / ( n J ) h ~ π − h ^ π \widetilde h_{\pi}-\widehat h_{\pi} h π − h π V n = O p ( 1 ) V_{n}=O_{p}(1) V n = O p ( 1 ) ∥ h ~ π − h ^ π ∥ H = O p ( ( n J ) − 1 / 2 ) \|\widetilde h_{\pi}-\widehat h_{\pi}\|_{\mathcal{H}}=O_{p}((nJ)^{-1/2}) ∥ h π − h π ∥ H = O p (( n J ) − 1/2 )
Proof of Corollary 5.6. Define
λ o p t : = τ κ 2 d Z c β 2 16 ( 1 − 2 α ) . \lambda_{\mathrm{opt}}:=\frac{\tau\kappa^{2}d_{Z}c_{\beta}^{2}}{16(1-2\alpha)}. λ opt := 16 ( 1 − 2 α ) τ κ 2 d Z c β 2 . By the definition of C R C_{R} C R λ o p t C R = s \lambda_{\mathrm{opt}}C_{R}=s λ opt C R = s
We first control the numerical target error. Since J ≥ 1 J \ge1 J ≥ 1 s ∈ ( 0 , 1 / 2 ] s \in(0,1/2] s ∈ ( 0 , 1/2 ] r m , n r w , n = O ( n − s ) r_{m,n}r_{w,n}=O(n^{-s}) r m , n r w , n = O ( n − s )
e n , J = O p ( n − 1 / 2 + r m , n r w , n + ( n J ) − 1 / 2 ) = O p ( n − s ) = o p ( 1 ) . e_{n,J}=O_p\left(n^{-1/2}+r_{m,n}r_{w,n}+(nJ)^{-1/2}\right)=O_p(n^{-s})=o_p(1). e n , J = O p ( n − 1/2 + r m , n r w , n + ( n J ) − 1/2 ) = O p ( n − s ) = o p ( 1 ) . We next verify the target-norm event required by Theorem 5.3 . Let
Δ κ : = κ − sup z ∈ Z ∥ φ ( z ) ∥ H > 0 \Delta_\kappa:=\kappa-\sup_{z\in\mathcal{Z}}\|\varphi(z)\|_{\mathcal{H}}>0 Δ κ := κ − z ∈ Z sup ∥ φ ( z ) ∥ H > 0 and define
E n : = { ∥ h ~ π ∥ H ≤ κ } . \mathcal{E}_n:=\left\{\|\widetilde{h}_\pi\|_{\mathcal{H}}\le\kappa\right\}. E n := { ∥ h π ∥ H ≤ κ } . Since
∥ χ ( π ) ∥ H = ∥ ∫ Z φ ( z ) Q π ( d z ) ∥ H ≤ sup z ∈ Z ∥ φ ( z ) ∥ H , \|\chi(\pi)\|_{\mathcal{H}}=\left\|\int_{\mathcal{Z}}\varphi(z)\,Q^\pi(dz)\right\|_{\mathcal{H}}\le\sup_{z\in\mathcal{Z}}\|\varphi(z)\|_{\mathcal{H}}, ∥ χ ( π ) ∥ H = ∫ Z φ ( z ) Q π ( d z ) H ≤ z ∈ Z sup ∥ φ ( z ) ∥ H , the triangle inequality gives
∥ h ~ π ∥ H ≤ sup z ∈ Z ∥ φ ( z ) ∥ H + e n , J . (125) \|\widetilde{h}_\pi\|_{\mathcal{H}}\le\sup_{z\in\mathcal{Z}}\|\varphi(z)\|_{\mathcal{H}}+e_{n,J}. \tag*{(125)} ∥ h π ∥ H ≤ z ∈ Z sup ∥ φ ( z ) ∥ H + e n , J . (125) Consequently,
Pr ( E n c ) ≤ Pr ( e n , J > Δ κ ) ⟶ 0. \Pr(\mathcal{E}_n^c)\le\Pr(e_{n,J}>\Delta_\kappa)\longrightarrow0. Pr ( E n c ) ≤ Pr ( e n , J > Δ κ ) ⟶ 0. Fix δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 ) E n \mathcal{E}_n E n Theorem 5.3 with R = R n R=R_n R = R n M = M n M=M_n M = M n (20) , with conditional probability at least 1 − δ 1-\delta 1 − δ
MMD k ( Q ^ M n ( R n ) , Q π ) ≤ C 0 exp { − λ o p t R n 1 − 2 α } + ( 1 + C 1 R n α ) e n , J + C 2 R n α log ( R n / δ ) M n . (126) \operatorname{MMD}_k\left(\widehat{Q}^{(R_n)}_{M_n},Q^\pi\right)\le C_0\exp\left\{-\lambda_{\mathrm{opt}}R_n^{1-2\alpha}\right\}+(1+C_1R_n^\alpha)e_{n,J}+C_2R_n^\alpha\sqrt{\frac{\log(R_n/\delta)}{M_n}}. \tag*{(126)} MMD k ( Q M n ( R n ) , Q π ) ≤ C 0 exp { − λ opt R n 1 − 2 α } + ( 1 + C 1 R n α ) e n , J + C 2 R n α M n log ( R n / δ ) . (126) We now control the three terms in Equation (126) . By Equation (23) ,
R n ≥ ( C R log n ) 1 / ( 1 − 2 α ) . R_n\ge(C_R\log n)^{1/(1-2\alpha)}. R n ≥ ( C R log n ) 1/ ( 1 − 2 α ) . Therefore,
exp { − λ o p t R n 1 − 2 α } ≤ exp { − λ o p t C R log n } = exp { − s log n } = n − s . \exp\left\{-\lambda_{\mathrm{opt}}R_n^{1-2\alpha}\right\}\le\exp\{-\lambda_{\mathrm{opt}}C_R\log n\}=\exp\{-s\log n\}=n^{-s}. exp { − λ opt R n 1 − 2 α } ≤ exp { − λ opt C R log n } = exp { − s log n } = n − s . Moreover, R n α = O ( ( log n ) α / ( 1 − 2 α ) ) R_n^\alpha=O\left((\log n)^{\alpha/(1-2\alpha)}\right) R n α = O ( ( log n ) α / ( 1 − 2 α ) ) e n , J = O p ( n − s ) e_{n,J}=O_p(n^{-s}) e n , J = O p ( n − s )
( 1 + C 1 R n α ) e n , J = O p ( n − s ( log n ) α / ( 1 − 2 α ) ) . (1+C_1R_n^\alpha)e_{n,J}=O_p\left(n^{-s}(\log n)^{\alpha/(1-2\alpha)}\right). ( 1 + C 1 R n α ) e n , J = O p ( n − s ( log n ) α / ( 1 − 2 α ) ) . For fixed δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 ) log ( R n / δ ) = O ( log log n ) \log(R_n/\delta)=O(\log\log n) log ( R n / δ ) = O ( log log n ) n n n
M n ≍ n 2 s log log n . M_n\asymp n^{2s}\log\log n. M n ≍ n 2 s log log n . It follows that
R n α log ( R n / δ ) M n = O ( ( log n ) α / ( 1 − 2 α ) log log n n 2 s log log n ) = O ( n − s ( log n ) α / ( 1 − 2 α ) ) . R_n^\alpha\sqrt{\frac{\log(R_n/\delta)}{M_n}}
=O\left((\log n)^{\alpha/(1-2\alpha)}\sqrt{\frac{\log\log n}{n^{2s}\log\log n}}\right)
=O\left(n^{-s}(\log n)^{\alpha/(1-2\alpha)}\right). R n α M n log ( R n / δ ) = O ( ( log n ) α / ( 1 − 2 α ) n 2 s log log n log log n ) = O ( n − s ( log n ) α / ( 1 − 2 α ) ) . To pass from the conditional statement to an unconditional one, let G n ( δ ) \mathcal{G}_n(\delta) G n ( δ ) (126) holds. Since E n \mathcal{E}_n E n
Pr ( G n ( δ ) c ) ≤ Pr ( E n c ) + E [ 1 E n Pr ( G n ( δ ) c ∣ D ) ] ≤ o ( 1 ) + δ . \Pr(\mathcal{G}_n(\delta)^c)\le\Pr(\mathcal{E}_n^c)+\mathbb{E}\left[\mathbf{1}_{\mathcal{E}_n}\Pr(\mathcal{G}_n(\delta)^c\mid\mathcal{D})\right]\le o(1)+\delta. Pr ( G n ( δ ) c ) ≤ Pr ( E n c ) + E [ 1 E n Pr ( G n ( δ ) c ∣ D ) ] ≤ o ( 1 ) + δ . Since δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 )
MMD k ( Q ^ M n ( R n ) , Q π ) = O p ( n − s ( log n ) α / ( 1 − 2 α ) ) . \operatorname{MMD}_k\left(\widehat{Q}^{(R_n)}_{M_n},Q^\pi\right)=O_p\left(n^{-s}(\log n)^{\alpha/(1-2\alpha)}\right). MMD k ( Q M n ( R n ) , Q π ) = O p ( n − s ( log n ) α / ( 1 − 2 α ) ) . In particular, the right-hand side converges to zero. If MMD k \operatorname{MMD}_k MMD k Q ^ M n ( R n ) ⇒ Q π \widehat{Q}^{(R_n)}_{M_n}\Rightarrow Q^\pi Q M n ( R n ) ⇒ Q π
Outcome-space transfer through the frozen encoder–decoder
We now prove Proposition 5.7 to isolate the error induced by the mapping of the structured outcomes into the encoder’s latent space.
Proof of Proposition 5.7. First, k E k_E k E m ∈ N m \in\mathbb{N} m ∈ N y 1 , … , y m ∈ Y y_1,\ldots,y_m \in\mathcal{Y} y 1 , … , y m ∈ Y c 1 , … , c m ∈ R c_1,\ldots,c_m \in\mathbb{R} c 1 , … , c m ∈ R
∑ i = 1 m ∑ ℓ = 1 m c i c ℓ k E ( y i , y ℓ ) = ∑ i = 1 m ∑ ℓ = 1 m c i c ℓ ⟨ φ ( E ( y i ) ) , φ ( E ( y ℓ ) ) ⟩ H k = ∥ ∑ i = 1 m c i φ ( E ( y i ) ) ∥ H k 2 ≥ 0. \sum_{i=1}^{m}\sum_{\ell=1}^{m} c_i c_\ell k_E(y_i,y_\ell)
= \sum_{i=1}^{m}\sum_{\ell=1}^{m} c_i c_\ell\langle\varphi(E(y_i)),\varphi(E(y_\ell))\rangle_{\mathcal{H}_k}
= \left\lVert\sum_{i=1}^{m} c_i\varphi(E(y_i))\right\rVert_{\mathcal{H}_k}^{2} \ge0. i = 1 ∑ m ℓ = 1 ∑ m c i c ℓ k E ( y i , y ℓ ) = i = 1 ∑ m ℓ = 1 ∑ m c i c ℓ ⟨ φ ( E ( y i )) , φ ( E ( y ℓ )) ⟩ H k = i = 1 ∑ m c i φ ( E ( y i )) H k 2 ≥ 0. Thus, y ↦ φ ( E ( y ) ) y \mapsto\varphi(E(y)) y ↦ φ ( E ( y )) k E k_E k E
Let Q 1 , Q 2 Q_1,Q_2 Q 1 , Q 2 Z \mathcal{Z} Z E Z ∼ Q j ∥ E ( D ( Z ) ) − Z ∥ 2 < ∞ \mathbb{E}_{Z\sim Q_j}\lVert E(D(Z))-Z\rVert_2 < \infty E Z ∼ Q j ∥ E ( D ( Z )) − Z ∥ 2 < ∞ j ∈ { 1 , 2 } j \in\{1,2\} j ∈ { 1 , 2 } T = E ∘ D T=E\circ D T = E ∘ D j ∈ { 1 , 2 } j \in\{1,2\} j ∈ { 1 , 2 }
μ j = E Z ∼ Q j [ φ ( Z ) ] and μ ~ j = E Z ∼ Q j [ φ ( T ( Z ) ) ] . \mu_j=\mathbb{E}_{Z\sim Q_j}[\varphi(Z)]
\qquad\text{and}\qquad
\widetilde{\mu}_j=\mathbb{E}_{Z\sim Q_j}[\varphi(T(Z))]. μ j = E Z ∼ Q j [ φ ( Z )] and μ j = E Z ∼ Q j [ φ ( T ( Z ))] . The definition of a pushforward measure implies that, for every integrable function f f f ∫ f ( u ) ( T # Q j ) ( d u ) = ∫ f ( T ( z ) ) Q j ( d z ) \int f(u)\,(T_{\#}Q_j)(du)=\int f(T(z))\,Q_j(dz) ∫ f ( u ) ( T # Q j ) ( d u ) = ∫ f ( T ( z )) Q j ( d z )
MMD k E ( D # Q 1 , D # Q 2 ) = ∥ μ ~ 1 − μ ~ 2 ∥ H k = MMD k ( T # Q 1 , T # Q 2 ) . (127) \operatorname{MMD}_{k_E}(D_{\#}Q_1,D_{\#}Q_2)
= \lVert\widetilde{\mu}_1-\widetilde{\mu}_2\rVert_{\mathcal{H}_k}
= \operatorname{MMD}_{k}(T_{\#}Q_1,T_{\#}Q_2). \tag*{(127)} MMD k E ( D # Q 1 , D # Q 2 ) = ∥ μ 1 − μ 2 ∥ H k = MMD k ( T # Q 1 , T # Q 2 ) . (127) This is an exact identity, not an inequality. Insert and subtract μ 1 \mu_1 μ 1 μ 2 \mu_2 μ 2
μ ~ 1 − μ ~ 2 = ( μ 1 − μ 2 ) + ( μ ~ 1 − μ 1 ) − ( μ ~ 2 − μ 2 ) . \widetilde{\mu}_1-\widetilde{\mu}_2
=(\mu_1-\mu_2)+(\widetilde{\mu}_1-\mu_1)-(\widetilde{\mu}_2-\mu_2). μ 1 − μ 2 = ( μ 1 − μ 2 ) + ( μ 1 − μ 1 ) − ( μ 2 − μ 2 ) . The triangle inequality in H k \mathcal{H}_k H k
∥ μ ~ 1 − μ ~ 2 ∥ H k ≤ ∥ μ 1 − μ 2 ∥ H k + ∑ j = 1 2 ∥ μ ~ j − μ j ∥ H k = MMD k ( Q 1 , Q 2 ) + ∑ j = 1 2 ∥ μ ~ j − μ j ∥ H k . (128) \lVert\widetilde{\mu}_1-\widetilde{\mu}_2\rVert_{\mathcal{H}_k}
\le\lVert\mu_1-\mu_2\rVert_{\mathcal{H}_k}
+\sum_{j=1}^{2}\lVert\widetilde{\mu}_j-\mu_j\rVert_{\mathcal{H}_k}
= \operatorname{MMD}_{k}(Q_1,Q_2)
+\sum_{j=1}^{2}\lVert\widetilde{\mu}_j-\mu_j\rVert_{\mathcal{H}_k}. \tag*{(128)} ∥ μ 1 − μ 2 ∥ H k ≤ ∥ μ 1 − μ 2 ∥ H k + j = 1 ∑ 2 ∥ μ j − μ j ∥ H k = MMD k ( Q 1 , Q 2 ) + j = 1 ∑ 2 ∥ μ j − μ j ∥ H k . (128) For each j j j
μ ~ j − μ j = E Z ∼ Q j [ φ ( T ( Z ) ) − φ ( Z ) ] . \widetilde{\mu}_j-\mu_j
=\mathbb{E}_{Z\sim Q_j}[\varphi(T(Z))-\varphi(Z)]. μ j − μ j = E Z ∼ Q j [ φ ( T ( Z )) − φ ( Z )] . The norm of a Bochner integral is at most the integral of the norm. Hence,
∥ μ ~ j − μ j ∥ H k ≤ E Z ∼ Q j ∥ φ ( T ( Z ) ) − φ ( Z ) ∥ H k ≤ L k E Z ∼ Q j ∥ T ( Z ) − Z ∥ 2 = L k E Z ∼ Q j ∥ E ( D ( Z ) ) − Z ∥ 2 , (129) \begin{aligned}
\lVert\widetilde{\mu}_j-\mu_j\rVert_{\mathcal{H}_k}
&\le\mathbb{E}_{Z\sim Q_j}\lVert\varphi(T(Z))-\varphi(Z)\rVert_{\mathcal{H}_k}\\
&\le L_k\mathbb{E}_{Z\sim Q_j}\lVert T(Z)-Z\rVert_2\\
&=L_k\mathbb{E}_{Z\sim Q_j}\lVert E(D(Z))-Z\rVert_2, \tag*{(129)}
\end{aligned} ∥ μ j − μ j ∥ H k ≤ E Z ∼ Q j ∥ φ ( T ( Z )) − φ ( Z ) ∥ H k ≤ L k E Z ∼ Q j ∥ T ( Z ) − Z ∥ 2 = L k E Z ∼ Q j ∥ E ( D ( Z )) − Z ∥ 2 , (129) where the second inequality is precisely the assumed Lipschitz continuity of the canonical feature map. Combining Equations (127) –(129) proves
MMD k E ( D # Q 1 , D # Q 2 ) ≤ MMD k ( Q 1 , Q 2 ) + L k ∑ j = 1 2 E Z ∼ Q j ∥ E ( D ( Z ) ) − Z ∥ 2 . (130) \operatorname{MMD}_{k_E}(D_{\#}Q_1,D_{\#}Q_2)
\le\operatorname{MMD}_{k}(Q_1,Q_2)
+L_k\sum_{j=1}^{2}\mathbb{E}_{Z\sim Q_j}\lVert E(D(Z))-Z\rVert_2. \tag*{(130)} MMD k E ( D # Q 1 , D # Q 2 ) ≤ MMD k ( Q 1 , Q 2 ) + L k j = 1 ∑ 2 E Z ∼ Q j ∥ E ( D ( Z )) − Z ∥ 2 . (130) It remains to prove the claim involving the actual interventional law. Because Q π = E # P π Q_\pi=E_{\#}P_\pi Q π = E # P π
MMD k E ( D # Q ^ , P π ) = ∥ E Z ∼ Q ^ [ φ ( T ( Z ) ) ] − E Y ∼ P π [ φ ( E ( Y ) ) ] ∥ H k = ∥ E Z ∼ Q ^ [ φ ( T ( Z ) ) ] − E Z ∼ Q π [ φ ( Z ) ] ∥ H k = MMD k ( T # Q ^ , Q π ) . \begin{aligned}
\operatorname{MMD}_{k_E}(D_{\#}\widehat{Q},P_\pi)
= \left\lVert\mathbb{E}_{Z\sim\widehat{Q}}[\varphi(T(Z))]-\mathbb{E}_{Y\sim P_\pi}[\varphi(E(Y))]\right\rVert_{\mathcal{H}_k}
= \left\lVert\mathbb{E}_{Z\sim\widehat{Q}}[\varphi(T(Z))]-\mathbb{E}_{Z\sim Q_\pi}[\varphi(Z)]\right\rVert_{\mathcal{H}_k} \\
&= \operatorname{MMD}_{k}(T_{\#}\widehat{Q},Q_{\pi}).
\end{aligned} MMD k E ( D # Q , P π ) = E Z ∼ Q [ φ ( T ( Z ))] − E Y ∼ P π [ φ ( E ( Y ))] H k = E Z ∼ Q [ φ ( T ( Z ))] − E Z ∼ Q π [ φ ( Z )] H k = MMD k ( T # Q , Q π ) . Insert and subtract E Z ∼ Q ^ [ φ ( Z ) ] \mathbb{E}_{Z\sim\widehat{Q}}[\varphi(Z)] E Z ∼ Q [ φ ( Z )] (129) , now with Q j = Q ^ Q_j=\widehat{Q} Q j = Q
MMD k E ( D # Q ^ , P π ) ≤ MMD k ( Q ^ , Q π ) + L k E Z ∼ Q ^ ∥ E ( D ( Z ) ) − Z ∥ 2 . \operatorname{MMD}_{k_E}(D_{\#}\widehat{Q},P_{\pi})\leq\operatorname{MMD}_{k}(\widehat{Q},Q_{\pi})+L_k\mathbb{E}_{Z\sim\widehat{Q}}\lVert E(D(Z))-Z\rVert_2. MMD k E ( D # Q , P π ) ≤ MMD k ( Q , Q π ) + L k E Z ∼ Q ∥ E ( D ( Z )) − Z ∥ 2 . This proves the claim.
Supplementary on Numerical Experiments
We first recall the RKHS objects and kernels used by IMMD-Flow. We then give the exact objective, normalization, optimizer, and learning-rate schedules used in our experiments.
RKHS and Kernels
Let U U U U \mathcal{U} U k U : U × U → R k_U:\mathcal{U}\times\mathcal{U}\to\mathbb{R} k U : U × U → R H U \mathcal{H}_U H U φ U ( u ) = k U ( u , ⋅ ) \varphi_U(u)=k_U(u,\cdot) φ U ( u ) = k U ( u , ⋅ ) f ( u ) = ⟨ f , φ U ( u ) ⟩ H U f(u)=\langle f,\varphi_U(u)\rangle_{\mathcal{H}_U} f ( u ) = ⟨ f , φ U ( u ) ⟩ H U f ∈ H U f\in\mathcal{H}_U f ∈ H U [51 , 52 ] . If E [ k U ( U , U ) ] < ∞ \mathbb{E}[k_U(U,U)]<\infty E [ k U ( U , U )] < ∞ P P P
μ P : = E U ∼ P [ φ U ( U ) ] ∈ H U , E U ∼ P [ f ( U ) ] = ⟨ f , μ P ⟩ H U . \mu_P:=\mathbb{E}_{U\sim P}[\varphi_U(U)]\in\mathcal{H}_U,\qquad\mathbb{E}_{U\sim P}[f(U)]=\langle f,\mu_P\rangle_{\mathcal{H}_U}. μ P := E U ∼ P [ φ U ( U )] ∈ H U , E U ∼ P [ f ( U )] = ⟨ f , μ P ⟩ H U . For two laws P , Q P,Q P , Q
MMD k U ( P , Q ) : = ∥ μ P − μ Q ∥ H U . \operatorname{MMD}_{k_U}(P,Q):=\lVert\mu_P-\mu_Q\rVert_{\mathcal{H}_U}. MMD k U ( P , Q ) := ∥ μ P − μ Q ∥ H U . A characteristic kernel makes P ↦ μ P P\mapsto\mu_P P ↦ μ P P = Q P=Q P = Q [53 ] . Likewise, whenever the conditional Bochner expectation exists,
m U ∣ V ( v ) : = E [ φ U ( U ) ∣ V = v ] ∈ H U m_{U\mid V}(v):=\mathbb{E}[\varphi_U(U)\mid V=v]\in\mathcal{H}_U m U ∣ V ( v ) := E [ φ U ( U ) ∣ V = v ] ∈ H U is the conditional mean embedding of P U ∣ V = v P_{U\mid V=v} P U ∣ V = v [54 , 55 , 56 ] . In our construction, U = Z = E ( Y ) U=Z=E(Y) U = Z = E ( Y ) V = S = ( A , X ) V=S=(A,X) V = S = ( A , X ) m U ∣ V m_{U\mid V} m U ∣ V m 0 m_0 m 0
Kernels. For z , z ′ ∈ R d Z z,z'\in\mathbb{R}^{d_Z} z , z ′ ∈ R d Z
k R B F ( z , z ′ ; ℓ ) = exp ( − ∥ z − z ′ ∥ 2 2 2 ℓ 2 ) , ℓ > 0 , k_{\mathrm{RBF}}(z,z';\ell)=\exp\left(-\frac{\lVert z-z'\rVert_2^2}{2\ell^2}\right),\qquad\ell>0, k RBF ( z , z ′ ; ℓ ) = exp ( − 2 ℓ 2 ∥ z − z ′ ∥ 2 2 ) , ℓ > 0 , k I M Q ( z , z ′ ; c , β ) = ( c 2 + ∥ z − z ′ ∥ 2 2 ) β , c > 0 , β < 0. k_{\mathrm{IMQ}}(z,z';c,\beta)=\left(c^2+\lVert z-z'\rVert_2^2\right)^{\beta},\qquad c>0,\quad\beta<0. k IMQ ( z , z ′ ; c , β ) = ( c 2 + ∥ z − z ′ ∥ 2 2 ) β , c > 0 , β < 0. Both are positive-definite kernels. For the Riesz interaction, define
ρ s ( z , z ′ ) : = ∥ z − z ′ ∥ 2 s , κ s ( z , z ′ ) : = − ρ s ( z , z ′ ) , 0 < s < 2. \rho_s(z,z'):=\lVert z-z'\rVert_2^s,\qquad\kappa_s(z,z'):=-\rho_s(z,z'),\qquad0<s<2. ρ s ( z , z ′ ) := ∥ z − z ′ ∥ 2 s , κ s ( z , z ′ ) := − ρ s ( z , z ′ ) , 0 < s < 2. The function κ s \kappa_s κ s z 1 , … , z N z_1,\ldots,z_N z 1 , … , z N c 1 , … , c N c_1,\ldots,c_N c 1 , … , c N ∑ a = 1 N c a = 0 \sum_{a=1}^{N}c_a=0 ∑ a = 1 N c a = 0
∑ a = 1 N ∑ b = 1 N c a c b κ s ( z a , z b ) ≥ 0. \sum_{a=1}^{N}\sum_{b=1}^{N}c_ac_b\kappa_s(z_a,z_b)\geq0. a = 1 ∑ N b = 1 ∑ N c a c b κ s ( z a , z b ) ≥ 0. For any anchor z ⋆ z_{\star} z ⋆
k s , z ⋆ ( z , z ′ ) : = 1 2 { ρ s ( z , z ⋆ ) + ρ s ( z ′ , z ⋆ ) − ρ s ( z , z ′ ) } k_{s,z_{\star}}(z,z'):=\frac{1}{2}\left\{\rho_s(z,z_{\star})+\rho_s(z',z_{\star})-\rho_s(z,z')\right\} k s , z ⋆ ( z , z ′ ) := 2 1 { ρ s ( z , z ⋆ ) + ρ s ( z ′ , z ⋆ ) − ρ s ( z , z ′ ) } is positive definite. If X , X ′ ∼ i i d P X,X'\mathrel{\overset{\mathrm{iid}}{\sim}}P X , X ′ ∼ iid P Y , Y ′ ∼ i i d Q Y,Y'\mathrel{\overset{\mathrm{iid}}{\sim}}Q Y , Y ′ ∼ iid Q s s s
D s ( P , Q ) : = 2 E ρ s ( X , Y ) − E ρ s ( X , X ′ ) − E ρ s ( Y , Y ′ ) = 2 MMD k s , z ⋆ 2 ( P , Q ) . D_s(P,Q):=2\mathbb{E}\rho_s(X,Y)-\mathbb{E}\rho_s(X,X')-\mathbb{E}\rho_s(Y,Y')=2\operatorname{MMD}_{k_{s,z_{\star}}}^{2}(P,Q). D s ( P , Q ) := 2 E ρ s ( X , Y ) − E ρ s ( X , X ′ ) − E ρ s ( Y , Y ′ ) = 2 MMD k s , z ⋆ 2 ( P , Q ) . The anchor cancels from this discrepancy. Thus, following the Riesz flow literature, we use a Riesz kernel for the conditionally positive-definite interaction κ s \kappa_s κ s [25 , 57 ] .
Riesz Particle Optimization with AdamW and Warm-up Learning Rate
All reported Riesz experiments use s = 1 s=1 s = 1 G 1 , … , G M ∈ R d Z G_1,\ldots,G_M \in\mathbb{R}^{d_Z} G 1 , … , G M ∈ R d Z Z ˉ i \bar{Z}_i Z ˉ i i = 1 , … , n i=1,\ldots,n i = 1 , … , n
Z ˉ i = { E ( Y i ) − μ ˉ Z } ⊘ σ ˉ Z , \bar{Z}_i = \{E(Y_i)-\bar{\mu}_Z\}\oslash\bar{\sigma}_Z, Z ˉ i = { E ( Y i ) − μ ˉ Z } ⊘ σ ˉ Z , where μ ˉ Z \bar{\mu}_Z μ ˉ Z σ ˉ Z \bar{\sigma}_Z σ ˉ Z 10 − 3 10^{-3} 1 0 − 3 ⊘ \oslash ⊘ v i = v D R , i v_i=v_{\mathrm{DR},i} v i = v DR , i
D s ( G 1 : M ; v , Z ˉ 1 : n ) = 2 M ∑ j = 1 M ∑ i = 1 n v i ∥ G j − Z ˉ i ∥ 2 s − 1 M 2 ∑ j = 1 M ∑ ℓ = 1 M ∥ G j − G ℓ ∥ 2 s − ∑ i = 1 n ∑ q = 1 n v i v q ∥ Z ˉ i − Z ˉ q ∥ 2 s . (131) D_s(G_{1:M};v,\bar{Z}_{1:n}) = \frac{2}{M}\sum_{j=1}^{M}\sum_{i=1}^{n}v_i\lVert G_j-\bar{Z}_i\rVert_2^s-\frac{1}{M^2}\sum_{j=1}^{M}\sum_{\ell=1}^{M}\lVert G_j-G_\ell\rVert_2^s-\sum_{i=1}^{n}\sum_{q=1}^{n}v_iv_q\lVert\bar{Z}_i-\bar{Z}_q\rVert_2^s. \tag*{(131)} D s ( G 1 : M ; v , Z ˉ 1 : n ) = M 2 j = 1 ∑ M i = 1 ∑ n v i ∥ G j − Z ˉ i ∥ 2 s − M 2 1 j = 1 ∑ M ℓ = 1 ∑ M ∥ G j − G ℓ ∥ 2 s − i = 1 ∑ n q = 1 ∑ n v i v q ∥ Z ˉ i − Z ˉ q ∥ 2 s . (131) The last term depends only on the fixed target. We therefore optimize
L s ( G 1 : M ) = 2 M ∑ j = 1 M ∑ i = 1 n v i ∥ G j − Z ˉ i ∥ 2 s − 1 M 2 ∑ j = 1 M ∑ ℓ = 1 M ∥ G j − G ℓ ∥ 2 s . (132) \mathcal{L}_s(G_{1:M}) = \frac{2}{M}\sum_{j=1}^{M}\sum_{i=1}^{n}v_i\lVert G_j-\bar{Z}_i\rVert_2^s-\frac{1}{M^2}\sum_{j=1}^{M}\sum_{\ell=1}^{M}\lVert G_j-G_\ell\rVert_2^s. \tag*{(132)} L s ( G 1 : M ) = M 2 j = 1 ∑ M i = 1 ∑ n v i ∥ G j − Z ˉ i ∥ 2 s − M 2 1 j = 1 ∑ M ℓ = 1 ∑ M ∥ G j − G ℓ ∥ 2 s . (132) The first term couples the particles to the fixed weighted target, whereas the second induces repulsion among particles. When v i ≥ 0 v_i\ge0 v i ≥ 0 ∑ i v i = 1 \sum_i v_i=1 ∑ i v i = 1
Adaptive scale and particle gradient. At iteration r r r
( U 1 ( r ) , … , U M + n ( r ) ) : = ( G 1 ( r ) , … , G M ( r ) , Z ˉ 1 , … , Z ˉ n ) (U_1^{(r)},\ldots,U_{M+n}^{(r)}) := (G_1^{(r)},\ldots,G_M^{(r)},\bar{Z}_1,\ldots,\bar{Z}_n) ( U 1 ( r ) , … , U M + n ( r ) ) := ( G 1 ( r ) , … , G M ( r ) , Z ˉ 1 , … , Z ˉ n ) and compute
a r : = max { 10 − 3 , 1 M ( M + n ) d Z ∑ j = 1 M ∑ q = 1 M + n ∥ G j ( r ) − U q ( r ) ∥ 2 } . (133) a_r := \max\left\{10^{-3},\frac{1}{M(M+n)\sqrt{d_Z}}\sum_{j=1}^{M}\sum_{q=1}^{M+n}\lVert G_j^{(r)}-U_q^{(r)}\rVert_2\right\}. \tag*{(133)} a r := max { 1 0 − 3 , M ( M + n ) d Z 1 j = 1 ∑ M q = 1 ∑ M + n ∥ G j ( r ) − U q ( r ) ∥ 2 } . (133) Both the particles and target are divided by a r a_r a r s = 1 s=1 s = 1
g j ( r ) = 1 a r { 2 M ∑ i = 1 n v i ψ ( G j ( r ) − Z ˉ i ) − 2 M 2 ∑ ℓ = 1 M ψ ( G j ( r ) − G ℓ ( r ) ) } , (134) g_j^{(r)}=\frac{1}{a_r}\left\{\frac{2}{M}\sum_{i=1}^{n}v_i\psi(G_j^{(r)}-\bar{Z}_i)-\frac{2}{M^2}\sum_{\ell=1}^{M}\psi(G_j^{(r)}-G_\ell^{(r)})\right\}, \tag*{(134)} g j ( r ) = a r 1 { M 2 i = 1 ∑ n v i ψ ( G j ( r ) − Z ˉ i ) − M 2 2 ℓ = 1 ∑ M ψ ( G j ( r ) − G ℓ ( r ) ) } , (134) where ψ ( u ) = u / ∥ u ∥ 2 \psi(u)=u/\lVert u\rVert_2 ψ ( u ) = u / ∥ u ∥ 2 u ≠ 0 u\ne0 u = 0 ψ ( 0 ) = 0 \psi(0)=0 ψ ( 0 ) = 0
Optimizer and learning-rate schedule. We optimize the particle locations using AdamW [58 , 59 ] , with zero weight decay, ϵ A = 10 − 8 \epsilon_A=10^{-8} ϵ A = 1 0 − 8 Table 4 .
Table 4. Riesz–AdamW settings used for the reported experiments. All rows use W = 100 W = 100 W = 100 M = 2,000 M = 2{,}000 M = 2 , 000
Let η min = 10 − 6 \eta_{\min}=10^{-6} η m i n = 1 0 − 6 η max \eta_{\max} η m a x W W W R R R
η r = η min + ( η max − η min ) r W , 0 ≤ r < W . (135) \eta_r=\eta_{\min}+(\eta_{\max}-\eta_{\min})\frac{r}{W},\qquad0\le r<W. \tag*{(135)} η r = η m i n + ( η m a x − η m i n ) W r , 0 ≤ r < W . (135) After warm-up, we set
η r = η max ω r , W ≤ r < R , (136) \eta_r = \eta_{\max}\omega_r,\qquad W \le r < R, \tag*{(136)} η r = η m a x ω r , W ≤ r < R , (136) where
ω r = { 1 , constant schedule , 1 2 [ 1 + cos ( π r − W R − W ) ] , cosine schedule . (137) \omega_r =
\begin{cases}
1, & \text{constant schedule},\\
\frac{1}{2}\left[1+\cos\left(\pi\frac{r-W}{R-W}\right)\right], & \text{cosine schedule}.
\end{cases} \tag*{(137)} ω r = { 1 , 2 1 [ 1 + cos ( π R − W r − W ) ] , constant schedule , cosine schedule . (137) The algorithm performs exactly R R R O ( M ( n + M ) d Z ) O(M(n+M)d_Z) O ( M ( n + M ) d Z ) O ( M ( n + M ) ) O(M(n+M)) O ( M ( n + M ))
Details on Data-Generating Processes
We evaluate IMMD-Flow on three main benchmarks. ColorMNIST and CelebA evaluate structured image outcomes under discrete interventions and are adapted from Dance et al. [8 ] . dSprites evaluates a stochastic intervention with a continuous action and is adapted from the setting of Zenati et al. [13 ] . We additionally use a two-dimensional null-effect mixture Dance et al. [8 ] in our ablation studies (see Appendix § 15 .§ 6 ). We describe the exact constructions used in our experiments below.
ColorMNIST. We adapt the foreground-color confounding benchmark of Dance et al. [8 ] that uses MNIST images [60 ] . We generate
X ∼ Unif [ 0 , 1 ] , X \sim\operatorname{Unif}[0,1], X ∼ Unif [ 0 , 1 ] , A ∣ X = x ∼ Bernoulli ( σ ℓ ( w ( x − 1 / 2 ) ) ) , A \mid X=x \sim\operatorname{Bernoulli}\left(\sigma_{\ell}\left(w(x-1/2)\right)\right), A ∣ X = x ∼ Bernoulli ( σ ℓ ( w ( x − 1/2 ) ) ) , S ∣ A = a ∼ P ^ d a , ( d 0 , d 1 ) = ( 1 , 6 ) . S \mid A=a \sim\widehat{P}_{d_a},\qquad(d_0,d_1)=(1,6). S ∣ A = a ∼ P d a , ( d 0 , d 1 ) = ( 1 , 6 ) . Here, P ^ d a \widehat{P}_{d_a} P d a d a d_a d a S ∈ [ 0 , 1 ] 28 × 28 S \in[0,1]^{28 \times28} S ∈ [ 0 , 1 ] 28 × 28
M ( S ) = σ ℓ ( κ ( S − τ ) ) , c f g ( x ) = ( x , α , 1 − x ) T , Y = M ( S ) ⊙ ( c f g ( X ) ⊗ S ) . M(S)=\sigma_{\ell}\left(\kappa(S-\tau)\right),\qquad c_{\mathrm{fg}}(x)=(x,\alpha,1-x)^{\mathsf{T}},\qquad Y=M(S)\odot\left(c_{\mathrm{fg}}(X)\otimes S\right). M ( S ) = σ ℓ ( κ ( S − τ ) ) , c fg ( x ) = ( x , α , 1 − x ) T , Y = M ( S ) ⊙ ( c fg ( X ) ⊗ S ) . The sigmoid and mask operations are pixelwise, and the background is black. Our default configuration uses ( w , κ , τ , α ) = ( 10 , 20 , 0.08 , 0.2 ) (w,\kappa,\tau,\alpha)=(10,20,0.08,0.2) ( w , κ , τ , α ) = ( 10 , 20 , 0.08 , 0.2 ) w ∈ { 1 , … , 10 } w \in\{1,\ldots,10\} w ∈ { 1 , … , 10 } P Y ( a ) P_{Y(a)} P Y ( a ) d a d_a d a X ∼ Unif [ 0 , 1 ] X \sim\operatorname{Unif}[0,1] X ∼ Unif [ 0 , 1 ] P ^ d a \widehat{P}_{d_a} P d a
CelebA. This benchmark evaluates population-level hair-color standardization within each sex group. We use the CelebA data set [61 ] and follow the controlled construction of Dance et al. [8 ] . Let
A = 1 { M a l e } , X = 1 { B l o n d H a i r } . A=1\{\mathrm{Male}\},\qquad X=1\{\mathrm{Blond\ Hair}\}. A = 1 { Male } , X = 1 { Blond Hair } . We combine the official training and validation splits, retain images with exactly one of the blond-hair and brown-hair labels, and exclude images labeled black hair, gray hair, bald, or wearing a hat. Let P ^ a , x \widehat{P}_{a,x} P a , x ( A = a , X = x ) (A=a,X=x) ( A = a , X = x )
The logged design is
P o b s ( X = 1 ) = q = 0.3 , P o b s ( A = 1 ∣ X = 0 ) = 0.6 , P o b s ( A = 1 ∣ X = 1 ) = 0.1 , Y ∣ ( A = a , X = x ) ∼ P ^ a , x . P_{\mathrm{obs}}(X=1)=q=0.3,\qquad P_{\mathrm{obs}}(A=1\mid X=0)=0.6,\qquad P_{\mathrm{obs}}(A=1\mid X=1)=0.1,\qquad Y\mid(A=a,X=x)\sim\widehat{P}_{a,x}. P obs ( X = 1 ) = q = 0.3 , P obs ( A = 1 ∣ X = 0 ) = 0.6 , P obs ( A = 1 ∣ X = 1 ) = 0.1 , Y ∣ ( A = a , X = x ) ∼ P a , x . Conditional on each cell, images are sampled without replacement. Let n a x = # { i : A i = a , X i = x } n_{ax} = \#\{i : A_i = a, X_i = x\} n a x = # { i : A i = a , X i = x } n o b s = 20,000 n_{\mathrm{obs}} = 20{,}000 n obs = 20 , 000 ( n 00 , n 01 , n 10 , n 11 ) (n_{00}, n_{01}, n_{10}, n_{11}) ( n 00 , n 01 , n 10 , n 11 ) ( 5,600 , 5,400 , 8,400 , 600 ) (5{,}600, 5{,}400, 8{,}400, 600) ( 5 , 600 , 5 , 400 , 8 , 400 , 600 ) a a a
P a ∗ = ( 1 − q ) P ^ a , 0 + q P ^ a , 1 . P_a^* = (1-q)\widehat{P}_{a,0} + q\widehat{P}_{a,1}. P a ∗ = ( 1 − q ) P a , 0 + q P a , 1 . We use 2,000 held-out reference images per group: 1,400 brown-haired and 600 blond-haired images. Logged and reference images are disjoint within every cell. This is population-level post-stratification within a sex group; it is not a same-person edit or a literal intervention on sex. We center-crop the aligned images to 178 × 178 178 \times178 178 × 178 256 × 256 256 \times256 256 × 256 [ − 1 , 1 ] [-1,1] [ − 1 , 1 ]
dSprites. We construct a confounded continuous-action benchmark from dSprites [62 ] , extending the setup of Zenati et al. [13 ] . Draw X = ( X 1 , X 2 ) ∼ Unif ( [ 0 , 1 ] 2 ) X = (X_1, X_2) \sim\operatorname{Unif}([0,1]^2) X = ( X 1 , X 2 ) ∼ Unif ([ 0 , 1 ] 2 )
C ( X ) = min { 2 , ⌊ 3 X 1 ⌋ } , S ( X ) = min { 5 , ⌊ 6 X 2 ⌋ } , C(X) = \min\{2, \lfloor3X_1\rfloor\}, \qquad S(X) = \min\{5, \lfloor6X_2\rfloor\}, C ( X ) = min { 2 , ⌊ 3 X 1 ⌋} , S ( X ) = min { 5 , ⌊ 6 X 2 ⌋} , which select one of three shapes and six scales; orientation is fixed at zero. Let R θ R_\theta R θ μ 0 ( x ) = 1.8 R π / 6 ( x − 1 2 1 2 ) \mu_0(x) = 1.8R_{\pi/6}(x - \frac{1}{2}\mathbf{1}_2) μ 0 ( x ) = 1.8 R π /6 ( x − 2 1 1 2 )
π 0 ( ⋅ ∣ x ) = N ( μ 0 ( x ) , 0.6 2 I 2 ) , \pi_0(\cdot\mid x) = \mathcal{N}(\mu_0(x), 0.6^2 I_2), π 0 ( ⋅ ∣ x ) = N ( μ 0 ( x ) , 0. 6 2 I 2 ) , π ( ⋅ ∣ x ) = 1 2 N ( μ 0 ( x ) , 0.6 2 I 2 ) + 1 2 N ( μ 0 ( x ) + Δ , 0.6 2 I 2 ) , Δ = ( 0.45 , − 0.45 ) ⊤ . \pi(\cdot\mid x) = \frac{1}{2}\mathcal{N}(\mu_0(x), 0.6^2 I_2) + \frac{1}{2}\mathcal{N}(\mu_0(x) + \Delta, 0.6^2 I_2), \qquad\Delta= (0.45,-0.45)^\top. π ( ⋅ ∣ x ) = 2 1 N ( μ 0 ( x ) , 0. 6 2 I 2 ) + 2 1 N ( μ 0 ( x ) + Δ , 0. 6 2 I 2 ) , Δ = ( 0.45 , − 0.45 ) ⊤ . The corresponding oracle importance ratio is
π ( a ∣ x ) π 0 ( a ∣ x ) = 1 2 + 1 2 exp { Δ ⊤ ( a − μ 0 ( x ) ) 0.6 2 − ∥ Δ ∥ 2 2 2 ( 0.6 2 ) } . \frac{\pi(a \mid x)}{\pi_0(a \mid x)} = \frac{1}{2} + \frac{1}{2}\exp\left\{\frac{\Delta^\top(a-\mu_0(x))}{0.6^2} - \frac{\lVert\Delta\rVert_2^2}{2(0.6^2)}\right\}. π 0 ( a ∣ x ) π ( a ∣ x ) = 2 1 + 2 1 exp { 0. 6 2 Δ ⊤ ( a − μ 0 ( x )) − 2 ( 0. 6 2 ) ∥ Δ ∥ 2 2 } . For j ∈ { 1 , 2 } j \in\{1,2\} j ∈ { 1 , 2 } p j ( a ) = 31 σ ℓ ( 1.5 a j ) p_j(a) = 31\sigma_\ell(1.5a_j) p j ( a ) = 31 σ ℓ ( 1.5 a j ) Y = g ( X , A ) ∈ [ 0 , 1 ] 64 × 64 Y = g(X,A) \in[0,1]^{64 \times64} Y = g ( X , A ) ∈ [ 0 , 1 ] 64 × 64 p ( A ) p(A) p ( A ) X X X X X X A A A A ← X → Y A \leftarrow X \to Y A ← X → Y X ⋆ ∼ P X X^\star\sim P_X X ⋆ ∼ P X A ⋆ ∼ π ( ⋅ ∣ X ⋆ ) A^\star\sim\pi(\cdot\mid X^\star) A ⋆ ∼ π ( ⋅ ∣ X ⋆ ) Y ⋆ = g ( X ⋆ , A ⋆ ) Y^\star= g(X^\star,A^\star) Y ⋆ = g ( X ⋆ , A ⋆ )
Two-dimensional Gaussian mixture. We use the null-effect deconfounding benchmark of Dance et al. [8 ] . Define σ ℓ ( t ) = ( 1 + e − t ) − 1 \sigma_\ell(t) = (1 + e^{-t})^{-1} σ ℓ ( t ) = ( 1 + e − t ) − 1 b = ( 1 , − 1 ) ⊤ / 2 b = (1,-1)^\top/\sqrt{2} b = ( 1 , − 1 ) ⊤ / 2
X ∼ N ( 0 , I 2 ) , A ∣ X = x ∼ Bernoulli ( σ ℓ ( γ b ⊤ x ) ) , X \sim\mathcal{N}(0,I_2), \qquad A \mid X = x \sim\operatorname{Bernoulli}(\sigma_\ell(\gamma b^\top x)), X ∼ N ( 0 , I 2 ) , A ∣ X = x ∼ Bernoulli ( σ ℓ ( γ b ⊤ x )) , B ∼ Bernoulli ( 1 / 2 ) , ϵ ∼ N ( 0 , σ Y 2 I 2 ) , Y = X + δ B 1 2 + ϵ . B \sim\operatorname{Bernoulli}(1/2), \qquad\epsilon\sim\mathcal{N}(0,\sigma_Y^2 I_2), \qquad Y = X + \delta B\mathbf{1}_2 + \epsilon. B ∼ Bernoulli ( 1/2 ) , ϵ ∼ N ( 0 , σ Y 2 I 2 ) , Y = X + δ B 1 2 + ϵ . Here, B B B ϵ \epsilon ϵ ( X , A ) (X,A) ( X , A ) ( γ , δ , σ Y ) = ( 3 , 2 , 0.2 ) (\gamma,\delta,\sigma_Y) = (3,2,0.2) ( γ , δ , σ Y ) = ( 3 , 2 , 0.2 ) P Y ( 0 ) = P Y ( 1 ) = P Y P_{Y(0)} = P_{Y(1)} = P_Y P Y ( 0 ) = P Y ( 1 ) = P Y P Y ∣ A = a P_{Y\mid A=a} P Y ∣ A = a P Y P_Y P Y
Autoencoders for Image Experiments
For ColorMNIST and dSprites, we train convolutional autoencoders on all logged images once per experimental repetition and freeze them before constructing the target embedding and running the particle flow. The ColorMNIST encoder maps a 3 × 28 × 28 3 \times28 \times28 3 × 28 × 28 1 × 64 × 64 1 \times64 \times64 1 × 64 × 64 [58 ] , with batch size 128 and learning rate 10 − 3 10^{-3} 1 0 − 3 § 5 treats E E E D D D
For CelebA, we instead use the frozen pretrained stabilityai/sd-vae-ft-mse VAE [63 ] , without finetuning, and encode each 256 × 256 256 \times256 256 × 256 4 × 32 × 32 4 \times32 \times32 4 × 32 × 32
Qualitative samples and particle evolution
In this section, we present qualitative examples of the initialized and generated particles. We also visualize selected particles at intermediate iterations to illustrate how the empirical distribution evolves under the fixed estimated target.
ColorMNIST samples. For an image y y y R ( y ) R(y) R ( y ) B ( y ) B(y) B ( y )
c ( y ) : = R ( y ) R ( y ) + B ( y ) . (138) c(y) := \frac{R(y)}{R(y) + B(y)}. \tag*{(138)} c ( y ) := R ( y ) + B ( y ) R ( y ) . (138) Values near zero correspond to predominantly blue foregrounds, values near one correspond to predominantly red foregrounds, and c ( y ) = 1 / 2 c(y) = 1/2 c ( y ) = 1/2 c ( Y ) = X c(Y) = X c ( Y ) = X X X X 1 − X 1 - X 1 − X X ∼ Unif [ 0 , 1 ] X \sim\operatorname{Unif}[0,1] X ∼ Unif [ 0 , 1 ] 1 / 2 1/2 1/2
Figure 3 compares particles before and after IMMD-Flow in both intervention arms. The initialized observational-arm particles have foreground-colour distributions shifted in opposite directions. The final particles move toward the interventional colour law while largely preserving the treatment-specific digit identity.
Figure 3. Initialized and final ColorMNIST particles. Within each pair of rows, columns correspond to the same numerical particles before and after optimization. The number displayed beside each row is the empirical mean of the foreground-colour statistic R / ( R + B ) R/(R+B) R / ( R + B )
CelebA samples. Figure 4 presents initialized and final particles for both sex groups. The generated particles remain visually coherent while changing under the fixed population-level hair-colour target.
Figure 4. Example CelebA particles before and after IMMD-Flow for the female and male groups.
Particle evolution on CelebA. Figure 5 shows selected particles at intermediate iterations. The decoded images change gradually rather than being replaced by unrelated samples, while the target embedding and its coefficients remain fixed throughout optimization.
Figure 5. Evolution of selected CelebA particles during IMMD-Flow. The top and bottom panels correspond to the female (a = 0 a = 0 a = 0 a = 1 a = 1 a = 1
Additional experiments and ablations
We complement the main benchmark comparisons with a low-dimensional illustration and a series of ablations examining the effects of the logged sample size, particle count, particle initialization, projected Nyström approximation, and confounding strength. Unless an experiment varies the stated component, we use the doubly robust target, full-sample nuisance fitting, projected Nyström regression, observed-outcome initialization, and M = 2,000 M = 2{,}000 M = 2 , 000
Entries are means and sample standard deviations across three data seeds. For the two static interventions, metrics are first averaged across intervention arms within each seed. For dSprites, they are averaged across ten fresh target-policy draws within each seed. Outcome- and latent-space S W 2 \mathrm{SW}_2 SW 2 10 3 10^3 1 0 3
Null-effect deconfounding illustration. We first illustrate IMMD-Flow on the two-dimensional null-effect benchmark of Dance et al. [8 ] which is explained in the previous section. Treatment does not enter the outcome equation, so
P Y ( 0 ) = P Y ( 1 ) = P Y . P_{Y(0)} = P_{Y(1)} = P_Y. P Y ( 0 ) = P Y ( 1 ) = P Y . Nevertheless, treatment assignment depends on X X X P Y ∣ A = a P_{Y \mid A=a} P Y ∣ A = a P Y P_Y P Y
Figure 6 compares particles initialized from each observational arm with the particles returned by DR–IMMD. The initial clouds overrepresent different regions of the common target, whereas the flowed particles recover its overall mixture structure more accurately. Figure 7 visualizes representative particle paths during optimization. The trajectories are generally curved because each particle follows a time-varying MMD witness determined by the evolving empirical particle distribution. They are not intended to approximate straight optimal-transport displacement paths.
Figure 6. Initial and flowed particles in the two-dimensional benchmark of Dance et al. [8 ] . Rows correspond to the intervention arms a ∈ { 0 , 1 } a \in\{0,1\} a ∈ { 0 , 1 } P Y ( a ) = P Y P_Y(a) = P_Y P Y ( a ) = P Y P Y ∣ A = a P_{Y \mid A=a} P Y ∣ A = a n = 20,000 n = 20{,}000 n = 20 , 000 M = 2,000 M = 2{,}000 M = 2 , 000
Figure 7. Paths of 80 representative particles for each intervention arm in the run from Figure 6 over 3,000 RBF–Euler updates. Path color indicates the iteration; open red markers and black markers denote initial and returned particles, respectively. Gray points show a fresh oracle draw from P Y P_Y P Y
Effect of the logged sample size. Table 5 shows that IMMD-Flow substantially reduces outcome-space and, where applicable, latent-space S W 2 \mathrm{SW}_2 SW 2 n n n
Table 5: Ablation over the number of logged observations, before and after particle optimization. “Initial” denotes the particle distribution before the flow: the corresponding observational arm for the discrete-treatment benchmarks and the logged outcome distribution for dSprites. “Flow” denotes the DR–IMMD output. Lower is better. A dash denotes an unreported metric; the two-dimensional experiment has neither a separate latent-space metric nor image FID. All entries are means ± \pm ± n n n
Table 5.
Effect of the particle count. Table 6 shows that increasing M M M S W 2 \mathrm{SW}_{2} SW 2 M = 500 M = 500 M = 500 M M M
Table 6. Ablation over the number of particles before and after IMMD-Flow. “Initial” and “flow” denote the metrics before and after particle optimization, respectively. Lower is better. Entries are means ± \pm ± M M M M M M
The final FID also decreases with M M M M M M M M M
Effect of particle initialization. We examine the sensitivity of IMMD-Flow to its initial particle distribution. We compare initialization from observed outcomes with Gaussian alternatives of increasing distance from the observed outcome geometry. The moment-matched Gaussian uses the empirical mean and full covariance. For CelebA, the standardized normal is N ( 0 , I ) \mathcal{N}(0, I) N ( 0 , I ) N ( 0 , I ) \mathcal{N}(0, I) N ( 0 , I )
Table 7. Particle-initialization ablation at n = 20000 n = 20000 n = 20000 M = 2 , 000 M = 2,000 M = 2 , 000
Observed-outcome initialization gives the strongest results on ColorMNIST and CelebA, indicating that starting on the learned image manifold can materially facilitate particle optimization. On CelebA, the full moment-matched Gaussian attains almost the same latent SW2 _2 2 2 _2 2
In the two-dimensional benchmark, the preferred initialization depends on the particle update, while dSprites is nearly insensitive to the three initializations considered. Overall, observational initialization is most beneficial for the more complex image distributions, rather than being uniformly optimal.
Full ridge versus projected Nyström. We compare the full-ridge and projected-Nyström constructions of the conditional mean embedding on ColorMNIST. The paired runs use the same data seeds, intervention arms, and particle-flow settings. The projected estimator uses m L = 1024 m_L = 1024 m L = 1024 Table 8 reports the resulting outcome-space discrepancy and the wall-clock time required to solve the conditional-mean-embedding regression.
Table 8. ColorMNIST full-ridge and projected-Nyström comparison (M = 2,000 M = 2{,}000 M = 2 , 000 ± \pm ± × \times ×
The full-ridge solve scales cubically in n n n 0.20 → 1.26 → 9.09 0.20 \to1.26 \to9.09 0.20 → 1.26 → 9.09 0.03 0.03 0.03 0.04 0.04 0.04 6.7 × 6.7\times 6.7 × 227 × 227\times 227 ×
Effect of confounding strength. For ColorMNIST, the parameter w w w w w w w ∈ { 1 , … , 10 } w \in\{1, \ldots, 10\} w ∈ { 1 , … , 10 } n = 20,000 n = 20{,}000 n = 20 , 000 M = 2,000 M = 2{,}000 M = 2 , 000
Table 9. ColorMNIST ablation over confounding strength for DR–IMMD. “Initial” denotes the observational-arm particles before optimization, and “flow” denotes the particles returned by IMMD-Flow. Lower is better. Entries are means ± \pm ±
As w w w S W 2 \mathrm{SW}_2 SW 2
Compute resources
All experiments were run on an internal SLURM-managed cluster of Linux x86_64 nodes, using Python 3.12.12 with PyTorch 2.10.0 built against CUDA 12.6 (NVIDIA driver 580.95.05). Each run was submitted as a single-node, single-task job requesting one GPU, two to eight CPU cores, between 12 GB and 64 GB of system memory, and a wall-clock limit of four to twelve hours; parameter sweeps were executed as SLURM job arrays with one configuration per array task. The ColorMNIST and dSprites experiments used a single NVIDIA L40S GPU with approximately 46 GB of memory, and the CelebA experiments a single NVIDIA A100-SXM4-40GB with 40 GB; the two-dimensional synthetic design contains no GPU code and was run on CPU. Every wall-clock comparison we report was pinned to one GPU model, so no timing is compared across hardware.
Existing assets, licences, and sensitive content
All datasets used in this paper are existing research benchmarks and are used solely for research evaluation. ColorMNIST is derived from MNIST, which is made available under the Creative Commons Attribution–ShareAlike 3.0 Unported (CC BY-SA 3.0) licence. dSprites is obtained from the official Google DeepMind repository1 stabilityai/sd-vae-ft-mse checkpoint through the Hugging Face Diffusers implementation of AutoencoderKL.2
CelebA contains identifiable photographs of public figures together with binary appearance annotations, including the dataset’s sex and hair-colour labels. We use CelebA only as a controlled benchmark for population-level distributional rebalancing: we do not perform identity recognition or make person-specific causal claims. The results should not be interpreted as establishing that the generated distributions are demographically representative, fair, or appropriate for real-world use. No new human annotations or participant data were collected.
