---
title: 'AYF-LMD: Lagrangian Map Distillation'
url: https://www.emergentmind.com/topics/lagrangian-map-distillation-ayf-lmd
type: topic
---

# AYF-LMD: Lagrangian Map Distillation

Lagrangian Map Distillation (AYF-LMD) is a continuous-time generative modeling distillation framework designed to produce neural flow maps that efficiently and accurately approximate probability path dynamics between noise and data distributions. Developed within the Align Your Flow (AYF) methodology, AYF-LMD leverages the Lagrangian perspective to train two-time neural operators, enabling high sample quality in few sampling steps for both unconditional and conditional tasks, including high-resolution image and text-to-image synthesis [2506.14603][2406.07507].

## 1. Flow Maps and Two-Time Operators

AYF-LMD builds upon the formulation of two-time flow maps associated with time-dependent velocity fields. For a dynamical probabilistic process with state $x_t$ at time $t$ evolving according to $\frac{dx_t}{dt} = v_\phi(x_t, t)$, the flow map $f_\theta(x_t, t, s)$ aims to transport $x_t$ at time $t$ directly to the solution $x_s$ at any time $s$, respecting $f_\theta(x_t, t, t) = x_t$. In practice, $f_\theta$ is parameterized as
$$
f_\theta(x_t, t, s) = x_t + (s-t)\cdot F_\theta(x_t, t, s),
$$
where $F_\theta$ encodes the learned average velocity over $[t, s]$ [2506.14603].

## 2. Lagrangian Map Distillation Objective

The Lagrangian Map Distillation (LMD) loss enforces the neural flow map to satisfy the trajectory-level ODE:
$$
\frac{\partial}{\partial s} f_\theta(x_t, t, s) \approx v_\phi(f_\theta(x_t, t, s), s).
$$
The AYF-LMD objective is formalized as:
$$
\mathcal{L}_{\rm LMD}(\theta) = \mathbb{E}_{x_t, t, s}\left[ w(t,s) \left\| f_\theta(x_t, t, s) - \operatorname{ODE}_{s'\to s}[f_{\theta^-}(x_t, t, s')] \right\|_2^2 \right],
$$
where ODE$_{s'\to s}$ denotes a one-step Euler integration from $s'=s + \epsilon(t-s)$ to $s$ and the expectation is over the data and time distributions with weighting $w(t,s)$ [2506.14603]. 

Taking the infinitesimal limit ($\epsilon \rightarrow 0$) and differentiating yields the Lagrangian PINN loss [2406.07507]:
$$
\mathcal{L}_\mathrm{LMD}(\theta) = \iint w(s, t)\,\mathbb{E}_{x_s\sim\rho_s} \left| \frac{\partial}{\partial t} f_\theta(x_s, s, t) - v(f_\theta(x_s, s, t), t) \right|^2 ds\,dt,
$$
which matches the derivatives of the map with the learned velocity field over all time pairs. The loss admits analytic and empirical advantages: rapid convergence and stability for few-step distillation [2506.14603][2406.07507].

## 3. Unified Framework and Special Cases

AYF-LMD situates itself within the broader two-time map distillation framework, encompassing various fast generative modeling objectives as limiting cases:
- When $s \to t$, the LMD loss reduces to classical flow matching, which matches the instantaneous velocities.
- When applied with $s = 0$, other distillation objectives like continuous-time consistency models are recovered in the Eulerian formulation.
- The general two-time map LMD objective also unifies trajectory distillation and neural-operator regression losses: trajectory-based methods precompute true ODE solutions for distinct time pairs and enforce regression, while flow map matching and progressive distillation are interpreted as variations of learning approximate or composed two-time maps [2506.14603][2406.07507].

## 4. Training Methodology and Practical Implementation

AYF-LMD is realized via modern neural architectures (typically deep U-Nets), with time coordinates $s$ and $t$ sinusoidally embedded and fed throughout the network [2406.07507]. The map parameterization ensures the boundary condition $f_\theta(x, s, s) = x$ by design:
$$
\hat{X}_{s,t}(x) = (1-t+s)\cdot x + (t-s)\cdot f^\theta_{s,t}(x).
$$
The training loop involves:
- Sampling batches of initial states $x_s$ via interpolants between base and data distributions,
- Randomly drawing times $(s, t)$ (uniform or weighted sampling),
- Computing the map $y = f_\theta(x_s, s, t)$ and its time derivative with autodiff,
- Minimizing the squared error between the neural derivative and the target velocity field as specified above.

Typical optimization employs Adam with batch sizes in the $128$–$512$ regime, learning rates $\approx 10^{-4}$, and $10^5$–$3\cdot10^5$ iterations depending on dataset size and complexity. For teacher-student distillation (e.g. from large-parameterized velocity models), the trained teacher's vector field is used as the reference $v_\phi$ [2506.14603][2406.07507].

## 5. Performance and Empirical Behavior

AYF-LMD demonstrates robust empirical performance, yielding state-of-the-art few-step sample quality with substantially reduced inference cost. On benchmarks including ImageNet 64×64 and 512×512, AYF distilled models using AYF-LMD achieve class-conditional FID scores that outperform prior consistency and flow matching approaches at 1–8 sampling steps:
- On ImageNet 64×64, AYF-LMD achieves FID = 1.25 for $N_\mathrm{FE}=2$ steps and further improves with optional adversarial fine-tuning.
- On ImageNet 512×512, using a 280M-parameter model, AYF-LMD achieves FID = 1.87 at $N_\mathrm{FE}=2$ steps [2506.14603].

For text-to-image tasks, AYF-LMD distilled models were subjectively preferred by human raters over leading LoRA-based few-step samplers in user studies [2506.14603]. LMD converges more rapidly and stably than direct Flow Map Matching (FMM) losses, particularly in the few-step regime, as confirmed on CIFAR-10 and ImageNet-32 [2406.07507].

## 6. Extensions: Guidance and Adversarial Fine-tuning

Classical classifier-free guidance is prone to over-shooting at large scales; instead, AYF employs autoguidance, where a lower-quality checkpoint $v_\phi^\mathrm{weak}$ is linearly mixed with the velocity teacher $v_\phi$ via a scalar $\lambda$ sampled uniformly during training. All tangent calculations in distillation are adjusted accordingly:
$$
v_\phi^\mathrm{guided}(x_t, t) = \lambda v_\phi(x_t, t) + (1-\lambda) v_\phi^\mathrm{weak}(x_t, t),
$$
with $\lambda \in [1, 3]$ typical [2506.14603].

After AYF-LMD training, adversarial fine-tuning may be applied. A StyleGAN2 discriminator is used under a relativistic Softplus loss, regularized with R$_1$/R$_2$, and the total loss combines the LMD objective with adversarial feedback. This post-processing sharpens samples while preserving diversity and yields further improvements in FID across datasets [2506.14603].

## 7. Variants and Theoretical Foundations

AYF-LMD extends the Lagrangian flow map distillation paradigm introduced by Boffi et al. for consistency and flow-matching models, providing theoretical unification and extension. Related developments include:
- Initial/Terminal Velocity Matching (ITVM), which augments LMD with specialized matching terms at initial and terminal times using exponential moving average stabilization, leading to superior few-step performance in both low- and high-dimensional domains [2505.01169].
- Direct (velocity-free) flow-map training via stochastic interpolants is also possible, enabling self-consistent fitting of two-time maps without a pretrained velocity field. However, empirical results indicate that AYF-LMD with teacher guidance remains preferable for rapid convergence and top-tier few-step quality [2406.07507].

The Lagrangian map distillation framework provides a general, theoretically grounded strategy for scalable generative modeling distillation, supporting both data-driven and teacher-driven scenarios, and enabling superior trade-offs in sample quality, diversity, and computational efficiency.

Source: https://www.emergentmind.com/topics/lagrangian-map-distillation-ayf-lmd