Answers.Body: <p>Assume that <span class="math-container">$\sigma_t$</span> is sufficiently larger than the &quot;typical size&quot; of the samples of <span class="math-container">$p_0$</span>. Then we can make several heuristic observations:</p>
  <ul>
  <li><p><span class="math-container">$\mathbf{x} + \sigma_t \mathbf{z} \approx \sigma_t \mathbf{z}$</span> holds, since <span class="math-container">$\mathbf{x}$</span> is negligible compared to <span class="math-container">$\sigma_t \mathbf{z}$</span> with high probability.</p>
  </li>
  <li><p>On the other hand, the flow <span class="math-container">$f_t$</span> has the property that if <span class="math-container">$\mathbf{x}_t \sim p_t$</span>, then <span class="math-container">$f_t(\mathbf{x}_t) \sim p_0$</span>. Since <span class="math-container">$\mathbf{x} + \sigma_t \mathbf{z} \sim p_t$</span>, we have <span class="math-container">$f_t(\mathbf{x} + \sigma_t \mathbf{z}) \sim p_0$</span>.</p>
  </li>
  </ul>
  <p>From these two observations, we can expect that</p>
  <ol>
  <li><p><span class="math-container">$f_t(\mathbf{x} + \sigma_t \mathbf{z})$</span> is a sample of <span class="math-container">$p_0$</span> which is almost independent from <span class="math-container">$\mathbf{x}$</span> and is an almost deterministic function of <span class="math-container">$\mathbf{z}$</span>. Consequently,
  <span class="math-container">$$ p(f(\mathbf{x} + \sigma_t \mathbf{z}), \mathbf{z}) \approx p(f(\sigma_t \mathbf{z}), \mathbf{z}) $$</span></p>
  </li>
  <li><p>By the assumption, <span class="math-container">$f_t(\mathbf{x} + \sigma_t \mathbf{z})$</span> is also negligible compare to <span class="math-container">$\sigma_t \mathbf{z}$</span>. Consequently, <span class="math-container">$\tilde{\mathbf{x}}_t \approx \sigma_t \mathbf{z}$</span>.</p>
  </li>
  </ol>
  <hr />
  <p>As for the transportation cost, write <span class="math-container">$\mathbf{x}_t = \mathbf{x} + \sigma_t \mathbf{z}$</span> for simplicity and introduce the function</p>
  <p><span class="math-container">$$ I(t) = \mathbb{E}\bigl[ \| f_t(\mathbf{x}_t) - \mathbf{z}\|^2 \bigr], $$</span></p>
  <p>which measures the transportation cost using the square distance cost function. Note that</p>
  <p><span class="math-container">$$ I(0) = \mathbb{E}\bigl[ \| \mathbf{x} - \mathbf{z}\|^2 \bigr], $$</span></p>
  <p>hence OP's question can be rephrased to whether <span class="math-container">$ I(0) &gt; I(t)$</span> or not. For this, the following observation comes handy:</p>
  <blockquote>
  <p><strong>Lemma.</strong> In addition to OP's setting, assume that <span class="math-container">$v_t(\mathbf{x}) = \mathbb{E}[\dot{\mathbf{x}}_t \mid \mathbf{x}_t = \mathbf{x}]$</span>, which is the usual choice in the flow matching literature. Then <span class="math-container">$v_t$</span> generates the probability path <span class="math-container">$(p_t)$</span>, and
  <span class="math-container">\begin{align*}
  v_t(\mathbf{x})
  &amp;= \dot{\sigma}_t \mathbb{E}[z \mid \mathbf{x}_t = \mathbf{x}]
  = -\dot{\sigma}_t \sigma_t \nabla_\mathbf{x} \log p_t (\mathbf{x}).
  \end{align*}</span></p>
  </blockquote>
  <p>Using this, we can show that:</p>
  <blockquote>
  <p><strong>Claim.</strong> <span class="math-container">$I'(t) = -2\dot{\sigma}_t d + o(t)$</span> as <span class="math-container">$t \to 0^+$</span>. In particular, if <span class="math-container">$\sigma_t$</span> is continuously differentiable with <span class="math-container">$\dot{\sigma}_0 &gt; 0$</span>, then <span class="math-container">$I(t)$</span> is decreasing for small <span class="math-container">$t$</span>.</p>
  </blockquote>
  <p>I am not sure if this trend will continue as <span class="math-container">$t$</span> grows, but my hunch is that the cost will indeed remain smaller. This is indeed true in the following special case:</p>
  <blockquote>
  <p><strong>Special Case.</strong> Assume <span class="math-container">$d = 1$</span> and <span class="math-container">$x \sim \mathcal{N}(0, 1)$</span>. Then we can check that</p>
  <p><span class="math-container">$$ v_t(\mathbf{x}) = \frac{\dot{\sigma}_t \sigma_t}{1+\sigma_t^2} \mathbf{x}
  \qquad\text{and}\qquad
  f_t(\mathbf{x}) = \frac{1}{\sqrt{1+\sigma_t^2}} \mathbf{x}. $$</span></p>
  <p>Using this, we can find an analytic formula for <span class="math-container">$I(t)$</span> :</p>
  <p><span class="math-container">$$ I(t) = 2 \left(1 - \frac{\sigma_t}{\sqrt{1+\smash[b]{\sigma_t^2}}} \right) $$</span></p>
  <p>This is clearly decreasing in <span class="math-container">$t$</span>. Also, it is easy to check that <span class="math-container">$I'(0) = -2\dot{\sigma}_0$</span>, confirming the claim with <span class="math-container">$d = 1$</span>.</p>
  </blockquote>
  <hr />
  <p><em>Proof of Lemma.</em> For any test function <span class="math-container">$\varphi \in C^{\infty}_c(\mathbb{R}^d)$</span>, we evaluate the time-derivative of <span class="math-container">$\mathbb{E}[\varphi(\mathbf{x}_t)]$</span> in two ways. One one hand,</p>
  <p><span class="math-container">\begin{align*}
  \frac{\partial}{\partial t} \mathbb{E}[\varphi(\mathbf{x}_t)]
  &amp;= \frac{\partial}{\partial t} \int_{\mathbb{R}^d} \varphi(\mathbf{x}) p_t(\mathbf{x}) \, \mathrm{d}\mathbf{x}
  = \int_{\mathbb{R}^d} \varphi(\mathbf{x}) \frac{\partial}{\partial t} p_t(\mathbf{x}) \, \mathrm{d}\mathbf{x}.
  \end{align*}</span></p>
  <p>On the other hand, differentiating the random variable <span class="math-container">$\varphi(\mathbf{x}_t)$</span> directly and using the definition of <span class="math-container">$v_t$</span>,</p>
  <p><span class="math-container">\begin{align*}
  \frac{\partial}{\partial t} \mathbb{E}[\varphi(\mathbf{x}_t)]
  &amp;= \mathbb{E}[\nabla \varphi(\mathbf{x}_t) \cdot \dot{\mathbf{x}}_t] \\
  &amp;= \mathbb{E}[\nabla \varphi(\mathbf{x}_t) \cdot \mathbb{E}[\dot{\mathbf{x}}_t \mid \mathbf{x}_t]] \\
  &amp;= \mathbb{E}[\nabla \varphi(\mathbf{x}_t) \cdot v_t(\mathbf{x}_t)] \\
  &amp;= \int_{\mathbb{R}^d} \nabla \varphi(\mathbf{x}) \cdot v_t (\mathbf{x}) p_t(\mathbf{x}) \, \mathrm{d}\mathbf{x} \\
  &amp;= - \int_{\mathbb{R}^d} \varphi(\mathbf{x}) \nabla \cdot (v_t(\mathbf{x}) p_t(\mathbf{x})) \, \mathrm{d}\mathbf{x}
  \end{align*}</span></p>
  <p>where we used the law of iterated expectation in the second line and integration by parts in the last line. Now, since <span class="math-container">$\varphi$</span> is arbitrary, this shows that <span class="math-container">$v_t$</span> satisfies the continuity equation:</p>
  <p><span class="math-container">$$ \frac{\partial}{\partial t} p_t + \nabla \cdot (v_t p_t) = 0 $$</span></p>
  <p>It is well-known that this implies that <span class="math-container">$v_t$</span> generates the probability path <span class="math-container">$p_t$</span>. Next, invoking OP's settiing,</p>
  <p><span class="math-container">\begin{align*}
  v_t(\mathbf{x})
  &amp;= \mathbb{E}[ \dot{\mathbf{x}}_t \mid \mathbf{x}_t = \mathbf{x}]
  = \dot{\sigma}_t \mathbb{E}[\mathbf{z} \mid \mathbf{x}_t = \mathbf{x}] \\
  &amp;= \frac{\dot{\sigma}_t}{\sigma_t} \mathbb{E}[ \mathbf{x}_t - \mathbf{x}_0 \mid \mathbf{x}_t = \mathbf{x}]
  = \frac{\dot{\sigma}_t}{\sigma_t} (\mathbf{x} - \mathbb{E}[ \mathbf{x}_0 \mid \mathbf{x}_t = \mathbf{x}] ).
  \end{align*}</span></p>
  <p>Since <span class="math-container">$\mathbf{x}_0$</span> is independent of the gaussian variable <span class="math-container">$\mathbf{z}$</span> and <span class="math-container">$\mathbf{x}_t = \mathbf{x}_0 + \sigma_t \mathbf{z}$</span>, we can invoke Tweedie's formula to find:</p>
  <p><span class="math-container">$$ \mathbb{E}[ \mathbf{x}_0 \mid \mathbf{x}_t = \mathbf{x}] = \mathbf{x} + \sigma_t^2 \nabla \log p_t(\mathbf{x}). $$</span></p>
  <p>Plugging this back, we obtain the desired identity.</p>
  <hr />
  <p><em>Proof of Claim.</em> Note that <span class="math-container">$f_t^{-1}(\mathbf{y})$</span> sends the initial point <span class="math-container">$\mathbf{y}$</span> at time <span class="math-container">$0$</span> along the vector field <span class="math-container">$v_t$</span> up to time <span class="math-container">$t$</span>. Consequently, <span class="math-container">$f_t^{-1}$</span> is a flow with the corresponding vector field <span class="math-container">$v_t$</span>, i.e.,</p>
  <p><span class="math-container">$$\frac{\partial}{\partial t} f_t^{-1}(\mathbf{y}) = v_t(f_t^{-1}(\mathbf{y})).$$</span></p>
  <p>Using this observation, by differentiating both sides of the identity <span class="math-container">$\mathbf{x} = f_t(f_t^{-1}(\mathbf{x}))$</span> with respect to <span class="math-container">$t$</span>, we find that the material derivative of <span class="math-container">$f_t$</span> is zero:</p>
  <p><span class="math-container">$$ \frac{\partial f_t}{\partial t} + \frac{\partial f_t}{\partial \mathbf{x}} v_t = 0 $$</span></p>
  <p>Here, <span class="math-container">$\frac{\partial g}{\partial \mathbf{x}}$</span> stands for the <a href="https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant" rel="nofollow noreferrer">Jacobian matrix</a> of the multivariable function <span class="math-container">$g : \mathbb{R}^m \to \mathbb{R}^n$</span>, and all the vectors are regarded as column vectors. Combining these altogether,</p>
  <p><span class="math-container">\begin{align*}
  I'(t)
  &amp;= 2 \mathbb{E}\left[ \left\langle  \frac{\partial}{\partial t}(f_t(\mathbf{x}_t)), f_t(\mathbf{x}_t) - \mathbf{z} \right\rangle \right] \\
  &amp;= 2 \mathbb{E}\left[ \left\langle \frac{\partial f_t}{\partial t}(\mathbf{x}_t) + \frac{\partial f_t(\mathbf{x}_t)}{\partial \mathbf{x}_t} \dot{\mathbf{x}}_t, f_t(\mathbf{x}_t) - z \right\rangle \right] \\
  &amp;= 2 \mathbb{E}\left[ \left\langle \frac{\partial f_t(\mathbf{x}_t)}{\partial \mathbf{x}_t} \left( \dot{\mathbf{x}}_t - v_t(\mathbf{x}_t) \right), f_t(\mathbf{x}_t) - z \right\rangle \right] \\
  &amp;= 2 \dot{\sigma}_t \mathbb{E}\left[ \left\langle \frac{\partial f_t(\mathbf{x}_t)}{\partial \mathbf{x}_t} \left( \mathbf{z} - \mathbb{E}[\mathbf{z} \mid \mathbf{x}_t] \right), f_t(\mathbf{x}_t) - z \right\rangle \right]
  \end{align*}</span></p>
  <p>Now let <span class="math-container">$\mathbf{w}_t = \mathbf{z} - \mathbb{E}[\mathbf{z} \mid \mathbf{x}_t] $</span>. Then <span class="math-container">$\mathbb{E}[\mathbf{w}_t \mid \mathbf{x}_t] = 0$</span>, hence <span class="math-container">$\mathbb{E}[ \langle \mathbf{w}_t, g(\mathbf{x}_t) \rangle ] = 0$</span> for essentially any function <span class="math-container">$g : \mathbb{R}^d \to \mathbb{R}^d$</span>. Using this, we can further simplify the last line as:</p>
  <p><span class="math-container">\begin{align*}
  I'(t)
  &amp;= - 2 \dot{\sigma}_t \mathbb{E}\left[ \left\langle \frac{\partial f_t(\mathbf{x}_t)}{\partial \mathbf{x}_t} \mathbf{w}_t, \mathbf{z} \right\rangle \right] \\
  &amp;= - 2 \dot{\sigma}_t \mathbb{E}\left[ \left\langle \frac{\partial f_t(\mathbf{x}_t)}{\partial \mathbf{x}_t} \mathbf{w}_t, \mathbf{w}_t \right\rangle \right]
  \end{align*}</span></p>
  <p>Now, when <span class="math-container">$t = 0$</span>,</p>
  <ul>
  <li><p><span class="math-container">$f_0$</span> is the identity function, hence its Jacobian is the identity matrix: <span class="math-container">$\frac{\partial f_0}{\partial \mathbf{x}} = \mathbf{I}$</span>.</p>
  </li>
  <li><p><span class="math-container">$\mathbf{w}_0 = \mathbf{z} - \mathbb{E}[\mathbf{z} \mid x_0 = x] = \mathbf{z}$</span>, where the last equality follows from the independence between <span class="math-container">$\mathbf{z}$</span> and <span class="math-container">$\mathbf{x}$</span>.</p>
  </li>
  </ul>
  <p>Consequently,</p>
  <p><span class="math-container">$$ I'(0) = -2\dot{\sigma}_0 \mathbb{E}[\|\mathbf{z}\|^2] = -2\dot{\sigma}_0 d, $$</span></p>
  <p>and the conclusion follows from the continuity of <span class="math-container">$I'(t)$</span>.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-09-18T09:32:14.030Z
Answers.Text: Thanks a lot for this answer. This helps a lot. With colleagues, we have positive empirical results on using these $\tilde{x}_i$. Would you be interested in further collaborating on this project for a conference submission? If yes, how could I contact you?
Answers.UserDisplayName: tzb
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-09-18T11:00:18.770Z
Answers.Text: @tzb, Glad it helped! I'm also happy to hear that your idea is showing promising results. As for collaboration, I'm certainly interested and would be delighted if I can contribute further to your project. Regarding contact information, forgive me for redirecting you to [my CV](https://sites.google.com/view/sangchul-lee/portfolio/cv?authuser=0) where you can find my email address. I'm a bit scared of web crawlers! :)
Answers.UserDisplayName: Sangchul Lee
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-09-13T18:02:05.670Z
Answers.LastActivityDate: 2024-09-19T06:14:55.510Z
Answers.LastEditDate: 2024-09-19T06:14:55.510Z
Answers.UserDisplayName: Sangchul Lee
Body: <p>Let us assume <span class="math-container">$ p_0 $</span> is a probability distribution over <span class="math-container">$ \mathbb{R}^d $</span>. Let <span class="math-container">$ x_t $</span> be a diffusion process defined as:
  <span class="math-container">\begin{equation}
  x_t = x + \sigma_t z,
  \end{equation}</span>
  where <span class="math-container">$ x \sim p_0 $</span>, <span class="math-container">$ z \sim \mathit{N}(0,I) $</span>, and <span class="math-container">$ \sigma_t &lt; \sigma_{t+\varepsilon} $</span> for any <span class="math-container">$ \varepsilon &gt; 0 $</span>, with <span class="math-container">$ \sigma_t \in \mathbb{R}_+ $</span>. Thus, we have <span class="math-container">$ p(x_t) = p_0 * \mathit{N}(0, \sigma_t^2 I) $</span>, where <span class="math-container">$ * $</span> denotes the convolution operator.</p>
  <p>Now, assume that we have access to an ODE solver of the diffusion process:
  <span class="math-container">\begin{equation}
  \frac{dx}{dt} = v_t(x),
  \end{equation}</span></p>
  <p>with an associated flow function that satisfies:
  <span class="math-container">\begin{equation}
  f_t(x_t) = x_t - \int_0^t v_s(x_s) \, ds.
  \end{equation}</span></p>
  <p>Next, define the following random variable:</p>
  <p><span class="math-container">$$\tilde{x_t} = f_t(x + \sigma_t z) + \sigma_t z.$$</span></p>
  <p>What can be said about: (i) the distribution of <span class="math-container">$ \tilde{x_t} $</span>? (ii) the coupling between <span class="math-container">$ f_t(x + \sigma_t z) $</span> and <span class="math-container">$ z $</span>?</p>
  <p>My intuition is that the flow from the ODE solver, <span class="math-container">$ f_t(x + \sigma_t z) $</span>, enables better alignment (and reduced transport cost) between the pair <span class="math-container">$ (f_t(x + \sigma_t z), z) $</span> compared to <span class="math-container">$ (x, z) $</span>, but I was not able to formalize it.</p>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-09-13T10:56:52.590Z
Comments.Text: Does your $v_t$ have anything to do with your probability path $(p_t)_{t\in[0,1]}$? (For example, in flow matching we consider vector fields that generate the probability path.)
Comments.UserDisplayName: Sangchul Lee
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-09-13T11:04:56.257Z
Comments.Text: Also, how exactly is your $f_t$ defined? It seems to me that $f_t$ is not well-defined solely of the position $x_t$ at time $t$, requiring the information about the process all the way up to time $t$, i.e., $(x_s)_{s\in[0,t]}$. So, do you actually intend to consider the functional $$ f_t[(x_s)_{s\in[0,1]}] = x_t - \int_{0}^{t} v_s(x_s) \, \mathrm{d}s $$ acting on paths $(x_t)_{t\in[0,1]}$?
Comments.UserDisplayName: Sangchul Lee
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-09-13T12:30:05.000Z
Comments.Text: Indeed, this was not clear in my first post. My $v_t$ is indeed the vector field that generate the probability path. It allows to go backward in time from $x_t$ to $x_0$ by solving the ODE. $f_t$ is the flow map associated to the ODE, i.e. it corresponds to the mapping of $x_t$ to $x_0$ given by the ODE solver.
Comments.UserDisplayName: tzb
Tags: <diffusion-models><probability-theory>
Title: Characterization of a coupling based on an ODE solver
UserDisplayName: tzb
ThreadText: # Characterization of a coupling based on an ODE solver
  
  (Asked by: username_0 on 2024-09-13)
  
  Let us assume $ p\_0 $ is a probability distribution over $ \mathbb{R}^d $. Let $ x\_t $ be a diffusion process defined as:
  \begin{equation}
  x\_t = x + \sigma\_t z,
  \end{equation}
  where $ x \sim p\_0 $, $ z \sim \mathit{N}(0,I) $, and $ \sigma\_t < \sigma\_{t+\varepsilon} $ for any $ \varepsilon > 0 $, with $ \sigma\_t \in \mathbb{R}\_+ $. Thus, we have $ p(x\_t) = p\_0 \* \mathit{N}(0, \sigma\_t^2 I) $, where $ \* $ denotes the convolution operator.
  
  Now, assume that we have access to an ODE solver of the diffusion process:
  \begin{equation}
  \frac{dx}{dt} = v\_t(x),
  \end{equation}
  
  with an associated flow function that satisfies:
  \begin{equation}
  f\_t(x\_t) = x\_t - \int\_0^t v\_s(x\_s) \, ds.
  \end{equation}
  
  Next, define the following random variable:
  
  $$\tilde{x\_t} = f\_t(x + \sigma\_t z) + \sigma\_t z.$$
  
  What can be said about: (i) the distribution of $ \tilde{x\_t} $? (ii) the coupling between $ f\_t(x + \sigma\_t z) $ and $ z $?
  
  My intuition is that the flow from the ODE solver, $ f\_t(x + \sigma\_t z) $, enables better alignment (and reduced transport cost) between the pair $ (f\_t(x + \sigma\_t z), z) $ compared to $ (x, z) $, but I was not able to formalize it.
  
  **Tags:** diffusion-models, probability-theory
  > **Comments:**
  > **username_1**: Does your $v_t$ have anything to do with your probability path $(p_t)_{t\in[0,1]}$? (For example, in flow matching we consider vector fields that generate the probability path.)  
  > **username_1**: Also, how exactly is your $f_t$ defined? It seems to me that $f_t$ is not well-defined solely of the position $x_t$ at time $t$, requiring the information about the process all the way up to time $t$, i.e., $(x_s)_{s\in[0,t]}$. So, do you actually intend to consider the functional $$ f_t[(x_s)_{s\in[0,1]}] = x_t - \int_{0}^{t} v_s(x_s) \, \mathrm{d}s $$ acting on paths $(x_t)_{t\in[0,1]}$?  
  > **username_0**: Indeed, this was not clear in my first post. My $v_t$ is indeed the vector field that generate the probability path. It allows to go backward in time from $x_t$ to $x_0$ by solving the ODE. $f_t$ is the flow map associated to the ODE, i.e. it corresponds to the mapping of $x_t$ to $x_0$ given by the ODE solver.  
  
  ---
  
  ## 1 answers
  
  ### Answer by username_1 on 2024-09-13. Score: 2
  
  Assume that $\sigma\_t$ is sufficiently larger than the "typical size" of the samples of $p\_0$. Then we can make several heuristic observations:
  
  * $\mathbf{x} + \sigma\_t \mathbf{z} \approx \sigma\_t \mathbf{z}$ holds, since $\mathbf{x}$ is negligible compared to $\sigma\_t \mathbf{z}$ with high probability.
  * On the other hand, the flow $f\_t$ has the property that if $\mathbf{x}\_t \sim p\_t$, then $f\_t(\mathbf{x}\_t) \sim p\_0$. Since $\mathbf{x} + \sigma\_t \mathbf{z} \sim p\_t$, we have $f\_t(\mathbf{x} + \sigma\_t \mathbf{z}) \sim p\_0$.
  
  From these two observations, we can expect that
  
  1. $f\_t(\mathbf{x} + \sigma\_t \mathbf{z})$ is a sample of $p\_0$ which is almost independent from $\mathbf{x}$ and is an almost deterministic function of $\mathbf{z}$. Consequently,
     $$ p(f(\mathbf{x} + \sigma\_t \mathbf{z}), \mathbf{z}) \approx p(f(\sigma\_t \mathbf{z}), \mathbf{z}) $$
  2. By the assumption, $f\_t(\mathbf{x} + \sigma\_t \mathbf{z})$ is also negligible compare to $\sigma\_t \mathbf{z}$. Consequently, $\tilde{\mathbf{x}}\_t \approx \sigma\_t \mathbf{z}$.
  
  ---
  
  As for the transportation cost, write $\mathbf{x}\_t = \mathbf{x} + \sigma\_t \mathbf{z}$ for simplicity and introduce the function
  
  $$ I(t) = \mathbb{E}\bigl[ \| f\_t(\mathbf{x}\_t) - \mathbf{z}\|^2 \bigr], $$
  
  which measures the transportation cost using the square distance cost function. Note that
  
  $$ I(0) = \mathbb{E}\bigl[ \| \mathbf{x} - \mathbf{z}\|^2 \bigr], $$
  
  hence OP's question can be rephrased to whether $ I(0) > I(t)$ or not. For this, the following observation comes handy:
  
  > **Lemma.** In addition to OP's setting, assume that $v\_t(\mathbf{x}) = \mathbb{E}[\dot{\mathbf{x}}\_t \mid \mathbf{x}\_t = \mathbf{x}]$, which is the usual choice in the flow matching literature. Then $v\_t$ generates the probability path $(p\_t)$, and
  > \begin{align\*}
  > v\_t(\mathbf{x})
  > &= \dot{\sigma}\_t \mathbb{E}[z \mid \mathbf{x}\_t = \mathbf{x}]
  > = -\dot{\sigma}\_t \sigma\_t \nabla\_\mathbf{x} \log p\_t (\mathbf{x}).
  > \end{align\*}
  
  Using this, we can show that:
  
  > **Claim.** $I'(t) = -2\dot{\sigma}\_t d + o(t)$ as $t \to 0^+$. In particular, if $\sigma\_t$ is continuously differentiable with $\dot{\sigma}\_0 > 0$, then $I(t)$ is decreasing for small $t$.
  
  I am not sure if this trend will continue as $t$ grows, but my hunch is that the cost will indeed remain smaller. This is indeed true in the following special case:
  
  > **Special Case.** Assume $d = 1$ and $x \sim \mathcal{N}(0, 1)$. Then we can check that
  >
  > $$ v\_t(\mathbf{x}) = \frac{\dot{\sigma}\_t \sigma\_t}{1+\sigma\_t^2} \mathbf{x}
  > \qquad\text{and}\qquad
  > f\_t(\mathbf{x}) = \frac{1}{\sqrt{1+\sigma\_t^2}} \mathbf{x}. $$
  >
  > Using this, we can find an analytic formula for $I(t)$ :
  >
  > $$ I(t) = 2 \left(1 - \frac{\sigma\_t}{\sqrt{1+\smash[b]{\sigma\_t^2}}} \right) $$
  >
  > This is clearly decreasing in $t$. Also, it is easy to check that $I'(0) = -2\dot{\sigma}\_0$, confirming the claim with $d = 1$.
  
  ---
  
  *Proof of Lemma.* For any test function $\varphi \in C^{\infty}\_c(\mathbb{R}^d)$, we evaluate the time-derivative of $\mathbb{E}[\varphi(\mathbf{x}\_t)]$ in two ways. One one hand,
  
  \begin{align\*}
  \frac{\partial}{\partial t} \mathbb{E}[\varphi(\mathbf{x}\_t)]
  &= \frac{\partial}{\partial t} \int\_{\mathbb{R}^d} \varphi(\mathbf{x}) p\_t(\mathbf{x}) \, \mathrm{d}\mathbf{x}
  = \int\_{\mathbb{R}^d} \varphi(\mathbf{x}) \frac{\partial}{\partial t} p\_t(\mathbf{x}) \, \mathrm{d}\mathbf{x}.
  \end{align\*}
  
  On the other hand, differentiating the random variable $\varphi(\mathbf{x}\_t)$ directly and using the definition of $v\_t$,
  
  \begin{align\*}
  \frac{\partial}{\partial t} \mathbb{E}[\varphi(\mathbf{x}\_t)]
  &= \mathbb{E}[\nabla \varphi(\mathbf{x}\_t) \cdot \dot{\mathbf{x}}\_t] \\
  &= \mathbb{E}[\nabla \varphi(\mathbf{x}\_t) \cdot \mathbb{E}[\dot{\mathbf{x}}\_t \mid \mathbf{x}\_t]] \\
  &= \mathbb{E}[\nabla \varphi(\mathbf{x}\_t) \cdot v\_t(\mathbf{x}\_t)] \\
  &= \int\_{\mathbb{R}^d} \nabla \varphi(\mathbf{x}) \cdot v\_t (\mathbf{x}) p\_t(\mathbf{x}) \, \mathrm{d}\mathbf{x} \\
  &= - \int\_{\mathbb{R}^d} \varphi(\mathbf{x}) \nabla \cdot (v\_t(\mathbf{x}) p\_t(\mathbf{x})) \, \mathrm{d}\mathbf{x}
  \end{align\*}
  
  where we used the law of iterated expectation in the second line and integration by parts in the last line. Now, since $\varphi$ is arbitrary, this shows that $v\_t$ satisfies the continuity equation:
  
  $$ \frac{\partial}{\partial t} p\_t + \nabla \cdot (v\_t p\_t) = 0 $$
  
  It is well-known that this implies that $v\_t$ generates the probability path $p\_t$. Next, invoking OP's settiing,
  
  \begin{align\*}
  v\_t(\mathbf{x})
  &= \mathbb{E}[ \dot{\mathbf{x}}\_t \mid \mathbf{x}\_t = \mathbf{x}]
  = \dot{\sigma}\_t \mathbb{E}[\mathbf{z} \mid \mathbf{x}\_t = \mathbf{x}] \\
  &= \frac{\dot{\sigma}\_t}{\sigma\_t} \mathbb{E}[ \mathbf{x}\_t - \mathbf{x}\_0 \mid \mathbf{x}\_t = \mathbf{x}]
  = \frac{\dot{\sigma}\_t}{\sigma\_t} (\mathbf{x} - \mathbb{E}[ \mathbf{x}\_0 \mid \mathbf{x}\_t = \mathbf{x}] ).
  \end{align\*}
  
  Since $\mathbf{x}\_0$ is independent of the gaussian variable $\mathbf{z}$ and $\mathbf{x}\_t = \mathbf{x}\_0 + \sigma\_t \mathbf{z}$, we can invoke Tweedie's formula to find:
  
  $$ \mathbb{E}[ \mathbf{x}\_0 \mid \mathbf{x}\_t = \mathbf{x}] = \mathbf{x} + \sigma\_t^2 \nabla \log p\_t(\mathbf{x}). $$
  
  Plugging this back, we obtain the desired identity.
  
  ---
  
  *Proof of Claim.* Note that $f\_t^{-1}(\mathbf{y})$ sends the initial point $\mathbf{y}$ at time $0$ along the vector field $v\_t$ up to time $t$. Consequently, $f\_t^{-1}$ is a flow with the corresponding vector field $v\_t$, i.e.,
  
  $$\frac{\partial}{\partial t} f\_t^{-1}(\mathbf{y}) = v\_t(f\_t^{-1}(\mathbf{y})).$$
  
  Using this observation, by differentiating both sides of the identity $\mathbf{x} = f\_t(f\_t^{-1}(\mathbf{x}))$ with respect to $t$, we find that the material derivative of $f\_t$ is zero:
  
  $$ \frac{\partial f\_t}{\partial t} + \frac{\partial f\_t}{\partial \mathbf{x}} v\_t = 0 $$
  
  Here, $\frac{\partial g}{\partial \mathbf{x}}$ stands for the [Jacobian matrix](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant) of the multivariable function $g : \mathbb{R}^m \to \mathbb{R}^n$, and all the vectors are regarded as column vectors. Combining these altogether,
  
  \begin{align\*}
  I'(t)
  &= 2 \mathbb{E}\left[ \left\langle \frac{\partial}{\partial t}(f\_t(\mathbf{x}\_t)), f\_t(\mathbf{x}\_t) - \mathbf{z} \right\rangle \right] \\
  &= 2 \mathbb{E}\left[ \left\langle \frac{\partial f\_t}{\partial t}(\mathbf{x}\_t) + \frac{\partial f\_t(\mathbf{x}\_t)}{\partial \mathbf{x}\_t} \dot{\mathbf{x}}\_t, f\_t(\mathbf{x}\_t) - z \right\rangle \right] \\
  &= 2 \mathbb{E}\left[ \left\langle \frac{\partial f\_t(\mathbf{x}\_t)}{\partial \mathbf{x}\_t} \left( \dot{\mathbf{x}}\_t - v\_t(\mathbf{x}\_t) \right), f\_t(\mathbf{x}\_t) - z \right\rangle \right] \\
  &= 2 \dot{\sigma}\_t \mathbb{E}\left[ \left\langle \frac{\partial f\_t(\mathbf{x}\_t)}{\partial \mathbf{x}\_t} \left( \mathbf{z} - \mathbb{E}[\mathbf{z} \mid \mathbf{x}\_t] \right), f\_t(\mathbf{x}\_t) - z \right\rangle \right]
  \end{align\*}
  
  Now let $\mathbf{w}\_t = \mathbf{z} - \mathbb{E}[\mathbf{z} \mid \mathbf{x}\_t] $. Then $\mathbb{E}[\mathbf{w}\_t \mid \mathbf{x}\_t] = 0$, hence $\mathbb{E}[ \langle \mathbf{w}\_t, g(\mathbf{x}\_t) \rangle ] = 0$ for essentially any function $g : \mathbb{R}^d \to \mathbb{R}^d$. Using this, we can further simplify the last line as:
  
  \begin{align\*}
  I'(t)
  &= - 2 \dot{\sigma}\_t \mathbb{E}\left[ \left\langle \frac{\partial f\_t(\mathbf{x}\_t)}{\partial \mathbf{x}\_t} \mathbf{w}\_t, \mathbf{z} \right\rangle \right] \\
  &= - 2 \dot{\sigma}\_t \mathbb{E}\left[ \left\langle \frac{\partial f\_t(\mathbf{x}\_t)}{\partial \mathbf{x}\_t} \mathbf{w}\_t, \mathbf{w}\_t \right\rangle \right]
  \end{align\*}
  
  Now, when $t = 0$,
  
  * $f\_0$ is the identity function, hence its Jacobian is the identity matrix: $\frac{\partial f\_0}{\partial \mathbf{x}} = \mathbf{I}$.
  * $\mathbf{w}\_0 = \mathbf{z} - \mathbb{E}[\mathbf{z} \mid x\_0 = x] = \mathbf{z}$, where the last equality follows from the independence between $\mathbf{z}$ and $\mathbf{x}$.
  
  Consequently,
  
  $$ I'(0) = -2\dot{\sigma}\_0 \mathbb{E}[\|\mathbf{z}\|^2] = -2\dot{\sigma}\_0 d, $$
  
  and the conclusion follows from the continuity of $I'(t)$.
  
  > **Comments:**
  > **username_0**: Thanks a lot for this answer. This helps a lot. With colleagues, we have positive empirical results on using these $\tilde{x}_i$. Would you be interested in further collaborating on this project for a conference submission? If yes, how could I contact you?  
  > **username_1**: @tzb, Glad it helped! I'm also happy to hear that your idea is showing promising results. As for collaboration, I'm certainly interested and would be delighted if I can contribute further to your project. Regarding contact information, forgive me for redirecting you to [my CV](https://sites.google.com/view/sangchul-lee/portfolio/cv?authuser=0) where you can find my email address. I'm a bit scared of web crawlers! :)  
  
  ---
  
  