In-Distribution Finetuning (IDFT)

• In-Distribution Finetuning (IDFT) is a fine-tuning strategy that aligns optimization with a defined in-distribution reference to preserve pre-trained structures.
• It employs adaptive token weighting based on model-induced next-token distributions to balance in-distribution accuracy and out-of-distribution robustness.
• IDFT integrates techniques like feature-preserving adaptation and data reweighting, demonstrating performance gains (e.g., 5–10 pp improvement over standard fine-tuning).

Searching arXiv for the cited papers to ground the article in current records. In-Distribution Finetuning (IDFT) denotes a family of finetuning strategies that attempt to align optimization with an explicitly defined “in-distribution” reference rather than treating all training tokens, samples, or domains as equally reliable. In the most explicit formulation, introduced for LLM supervised fine-tuning, IDFT is a loss-level method derived from Distribution Discriminant Theory (DDT): it down-weights tokens that are unlikely under the model’s own next-token distribution and thereby seeks to improve generalization while preserving pre-trained structure (Zhang et al., 12 Feb 2026). In adjacent literatures, closely related ideas appear as target-distribution matching, feature-preserving finetuning, data reweighting, and quality-aligned adaptation, all motivated by the same empirical observation: naive full fine-tuning often improves in-distribution accuracy while degrading robustness or out-of-distribution behavior (Kumar et al., 2022).

1. Terminological scope and motivating problem

The literature uses the label “in-distribution” in more than one sense. In standard supervised fine-tuning (SFT) for LLMs, the relevant distribution is the model-induced next-token distribution $p_\theta(\cdot \mid c_t)$; IDFT then means suppressing gradients from tokens that behave as out-of-distribution relative to that distribution (Zhang et al., 12 Feb 2026). In transfer learning for vision, the relevant distribution is typically the target-domain data distribution $P_{in}$, and the problem is to adapt a pre-trained feature extractor $f_{\theta_0}$ and head $g_w$ while preserving the generalization of the pre-trained representation (Ha et al., 2023). In data-centric settings, the distributional question is whether the finetuning corpus matches the intended deployment distribution, whether defined by a small representative query set, by image-quality ratios, or by an inverse-density reweighting scheme (Liu et al., 2024).

The common motivation is that full fine-tuning can distort pre-trained features. One influential empirical study compared linear probing (LP), full fine-tuning (FT), and linear probing followed by full fine-tuning (LP-FT) across 10 distribution shift datasets and reported that fine-tuning obtained on average $2\%$ higher accuracy ID but $7\%$ lower accuracy OOD than linear probing; LP-FT was reported as $1\%$ better ID and $10\%$ better OOD than full fine-tuning (Kumar et al., 2022). This establishes the central tension that IDFT addresses across domains: optimization should adapt the model to the target task without overwriting the structures that support transfer and robustness.

2. Distribution Discriminant Theory and the canonical IDFT objective

The DDT formulation begins from the standard SFT loss

$$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$

with $c_t=(Q,x_{<t})$, vocabulary $P_{in}$0, and next-token distribution

$P_{in}$1

The paper contrasts this with an external data distribution $P_{in}$2 and frames “is this token in-distribution or out-of-distribution?” as the hypothesis test

$P_{in}$3

The statistic is optimized through a signal-to-noise ratio

$P_{in}$4

Within statistics of the form $P_{in}$5, Theorem 3.1 identifies the Centralized Log-Likelihood

$P_{in}$6

as the SNR-maximizing choice, where

$P_{in}$7

is the Shannon entropy (Zhang et al., 12 Feb 2026).

The sequential properties of $P_{in}$8 are central to the method. With $P_{in}$9, the cumulative statistic is a martingale under $f_{\theta_0}$0,

$f_{\theta_0}$1

while under $f_{\theta_0}$2 it has negative drift,

$f_{\theta_0}$3

A Freedman-inequality concentration bound is also given: $f_{\theta_0}$4 where $f_{\theta_0}$5 and $f_{\theta_0}$6. The practical implication is that strongly negative $f_{\theta_0}$7 values signal tokens that are misaligned with the model-induced distribution.

IDFT turns this statistic into an adaptive multiplier

$f_{\theta_0}$8

and defines the loss

$f_{\theta_0}$9

The paper states the role of the weighting explicitly: strongly OOD tokens with $g_w$0 yield $g_w$1 and $g_w$2, while high-confidence tokens with $g_w$3 are up-scaled. In implementation, $g_w$4 is clipped, for example with $g_w$5, before applying $g_w$6, and the resulting procedure adds only a small constant-factor overhead above SFT (Zhang et al., 12 Feb 2026).

3. Empirical profile in on-policy-style supervised finetuning

The LLM experiments position IDFT as a computationally efficient alternative to offline RL-style training. The base models are Phi-4-mini-instruct (3.8B) and Qwen2.5-7B-instruct; training uses 333K problems from Numina-Math; benchmarks are Math-C $g_w$7AIME24, AMC23, OlympiadBench$g_w$8, Math-G $g_w$9College-math, Math-OAI, Minerva-math$2\%$0, and General-R $2\%$1MMLU-stem, ARC-challenge$2\%$2. Baselines include SFT, DFT, EAFT, and offline RL methods such as Rej@2, Rej@4, DPO, SimPO, RPO, and CPO (Zhang et al., 12 Feb 2026).

The reported aggregate effect is that IDFT yields $2\%$3 pp relative gain over SFT and outperforms DFT and EAFT by $2\%$4 pp, thereby closing most of the gap to RL methods. A table excerpt for Phi4-mini gives the following values under the stated FLOPs budgets: Origin $2\%$5 / $2\%$6 / $2\%$7 on Math-C / Math-G / General-R; SFT $2\%$8 / $2\%$9 / $7\%$0; DPO $7\%$1 / $7\%$2 / $7\%$3; SimPO $7\%$4 / $7\%$5 / $7\%$6; HD+SFT $7\%$7 / $7\%$8 / $7\%$9; and HD+IDFT $1\%$0 / $1\%$1 / $1\%$2. Here HD denotes Hinted Decoding, the companion data-level technique that re-aligns the training corpus to the model’s distribution (Zhang et al., 12 Feb 2026).

The ablation evidence is also structurally important. Replacing the soft modulation $1\%$3 with a hard mask $1\%$4 produces an inverted-U curve as $1\%$5 varies, which the paper uses to argue that a continuous, self-tuning gate is superior to fixed truncation. This places IDFT in a broader class of adaptive-risk objectives: rather than rejecting entire examples, it modulates token contributions continuously according to model-relative in-distribution evidence.

4. Feature distortion, head alignment, and representation-preserving finetuning

A closely related line of work studies in-distribution adaptation through the lens of representation drift. In an overparameterized two-layer linear network

$1\%$6

full fine-tuning updates the head $1\%$7 and feature extractor $1\%$8 simultaneously. The analysis shows that when the head is initialized poorly, feature updates occur along directions spanned by the in-distribution data, leaving orthogonal directions almost unchanged; the resulting distorted features fit the training distribution but can perform poorly OOD. The same study reports that on 10 distribution shift benchmarks FT achieved $1\%$9 ID and $10\%$0 OOD accuracy on average, LP achieved $10\%$1 ID and $10\%$2 OOD, and LP-FT achieved $10\%$3 ID and $10\%$4 OOD (Kumar et al., 2022).

Domain-Aware Fine-Tuning (DAFT) operationalizes this concern in a target-domain setting with pre-trained feature extractor $10\%$5, head $10\%$6, and objective

$10\%$7

Its contribution is twofold: a batch-normalization conversion that replaces source-domain running statistics $10\%$8 with target-domain-adapted parameters $10\%$9, and an integrated LP+FT procedure with zero-initialized head and separate learning rates $$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$0. The converted BatchNorm parameters are

$$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$1

and feature distortion is quantified by

$$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$2

On the MoCo-v2 CIFAR-10 series, DAFT reports $$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$3 ID, $$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$4 on CIFAR-10.1, and $$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$5 on STL, compared with LP $$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$6, FT $$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$7, and LP-FT $$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$8; on CIFAR-10-C mean Corruption Error, lower is better, DAFT gives $$\mathcal{L}_{\rm SFT}(\theta) \;=\; -\frac{1}{L}\sum_{t=1}^L \log p_\theta(x_t\mid c_t),$$9 versus LP $c_t=(Q,x_{<t})$0, FT $c_t=(Q,x_{<t})$1, and LP-FT $c_t=(Q,x_{<t})$2 (Ha et al., 2023).

FD-Align addresses the same problem from a feature-alignment perspective. It is introduced as “Feature Discrimination Alignment,” a fine-tuning approach that aims to bolster generalizability by preserving the consistency of spurious features across the fine-tuning process. The available summary states that it decouples features into category-related and category-independent information, imposes an alignment constraint that penalizes changes in label-independent feature channels or subspaces, and demonstrates efficacy for both ID and OOD tasks; once fine-tuned, the model can seamlessly integrate with existing methods, leading to performance improvements (Song et al., 2023). Taken together, these methods suggest that a substantial portion of what is called IDFT in transfer learning is better understood as controlled adaptation: align the head, constrain the backbone, and preserve useful pre-trained structure.

5. Data-centric forms of IDFT: selection, quality matching, and deployment alignment

Another line of work treats IDFT as a data selection or data composition problem. TSDS (Task-Specific Data Selection) begins from a small representative set $c_t=(Q,x_{<t})$3 from the target task distribution and a large candidate repository $c_t=(Q,x_{<t})$4. It formulates data selection for task-specific model finetuning as an optimal-transport problem with a distribution alignment loss

$c_t=(Q,x_{<t})$5

followed by marginalization $c_t=(Q,x_{<t})$6. To discourage near-duplicates, TSDS adds a kernel-density-based regularizer with

$c_t=(Q,x_{<t})$7

yielding a KDE-aware transport objective. The framework is evaluated for continued pretraining and instruction tuning and is reported to show that instruction tuning using data selected with a $c_t=(Q,x_{<t})$8 selection ratio often outperforms using the full dataset and beats the baseline selection methods by $c_t=(Q,x_{<t})$9 points in F1 score on average (Liu et al., 2024).

In medical imaging, the same distribution-matching principle appears in the relation between finetuning and deployment quality distributions. For prostate multiparametric MRI, the study defines high-quality and low-quality exams through PI-QUAL, with $P_{in}$00 denoting the finetuning distribution and $P_{in}$01 the test distribution, where

$P_{in}$02

and

$P_{in}$03

The ideal in-distribution regime is $P_{in}$04. The experiments vary $P_{in}$05 and show systematic degradation under mismatch. For PI-RADS 5-class, the HQ-only model’s T1-ACC falls from approximately $P_{in}$06 at $P_{in}$07 to approximately $P_{in}$08 at $P_{in}$09, while the Mixed model declines more gracefully. Against training from scratch under identical conditions, HQ-only finetuning yields average ACC gains of $P_{in}$10 and AUC gains $P_{in}$11, Mixed finetuning yields $P_{in}$12 ACC gain, and LQ-only yields $P_{in}$13 ACC change. The study therefore recommends aligning quality distributions and notes that a minimum HQ proportion of approximately $P_{in}$14 in finetuning is advised for tasks where scratch accuracy gains are critical (Tang et al., 16 Aug 2025).

6. Theoretical extensions, misconceptions, and open questions

A broader theoretical perspective treats IDFT as gentle finetuning or rebalancing toward a more uniform domain coverage. One formulation defines

$P_{in}$15

where $P_{in}$16 is the uniform density on $P_{in}$17, and studies the worst-case risk over all sufficiently diverse test distributions. The main result is that training on the uniform distribution is minimax-optimal for this worst-case objective. When uniform samples are unavailable, the proposed operationalization is an IDFT objective

$P_{in}$18

with reweighting $P_{in}$19 and a small, decaying learning rate so that the fine-tuned model remains close to the pre-trained prior. A PAC-Bayes bound then controls the worst-case diverse-distribution risk through the divergence between posterior and prior over weights (Loukas et al., 2024).

Several misconceptions are corrected by this literature. First, IDFT is not a single algorithmic recipe. Depending on context, it may denote model-distribution-aware token weighting, inverse-density reweighting, head-aligned feature-preserving adaptation, or finetune/test distribution matching. Second, “better ID accuracy” is not equivalent to “better generalization”: full fine-tuning can outperform LP in-distribution while underperforming OOD, and medical-image quality mismatch can negate the gains from pretraining altogether (Kumar et al., 2022). Third, the relevant distribution need not be the empirical training distribution; in DDT-based IDFT, the reference is the model-induced distribution itself, and in quality-aware finetuning the critical variable is the mismatch between the finetuning and deployment distributions (Zhang et al., 12 Feb 2026).

The open questions are correspondingly heterogeneous. The LLM formulation highlights calibration sensitivity and the need to clip $P_{in}$20. The quality-aware medical setting explicitly points to adapter-based or quality-aware sampling strategies and richer low-quality datasets. The rebalancing framework raises the problem of density estimation in high dimensions and outside-support shift. Across these strands, however, the underlying principle remains stable: finetuning is most effective when its gradients, data weights, and sample composition are made commensurate with the distribution on which the adapted model is expected to operate (Tang et al., 16 Aug 2025).