feat(batch17): RoPE scaling extensions, DPO loss, label smoothing + focal loss

- RopeTable/RopeScaler: linear interpolation, dynamic NTK (base scaling),
  YaRN per-frequency blending with ramp fn + temperature correction
  (arXiv:2309.00071); apply_to_sequence multi-head; 23 tests
- DpoLoss: log-sigmoid DPO (arXiv:2305.18290), IPO squared variant
  (arXiv:2310.12036), robust DPO label smoothing; implicit reward tracking;
  DpoAccumulator with preference accuracy; 23 tests
- LossFunctions: label-smoothed CE (Szegedy 2016), focal loss (Lin 2017
  arXiv:1708.02002), smoothed focal, binary CE (stable), binary focal;
  Reduction::Mean/Sum/None; 22 tests

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

crates/training/rtx-transformers/src/layers/mod.rs

@@ -102,6 +102,7 @@ pub mod shared_layers;
 // Position encoding variants
 pub mod rope_cuda;        // Fused RoPE kernel — CPU reference + cudarc GPU dispatch
 pub mod rope_improvements; // DynamicRoPE, RoPE2D, InterpolatedRoPE, XPos, RoPECache
+pub mod rope_scaling;     // Linear / Dynamic-NTK / YaRN frequency scaling for long context
 // pub mod ntk_rope;
 // pub mod alibi;
 
@@ -240,6 +241,7 @@ pub use shared_layers::{SharedFfnWeight, SharedLayerConfig, SharedLayerStack, Sh
 pub use sliding_window_attn::{
     AttentionStats, SlidingWindowAttention, SlidingWindowConfig, WindowMask,
 };
+pub use rope_scaling::{RopeScalingConfig, RopeScalingMode, RopeTable, RopeScaler};
 
 // TransformerConfig is defined above and available for import
 

crates/training/rtx-transformers/src/layers/rope_scaling.rs

@@ -0,0 +1,988 @@
+//! RoPE frequency scaling extensions for long-context inference.
+//!
+//! Standard RoPE (Rotary Position Embedding) breaks when sequences exceed the
+//! training length because the model encounters position frequencies it never
+//! saw during training. This module implements three remedies:
+//!
+//! - **Linear interpolation** (Chen et al. 2023, arXiv:2306.15595): scales all
+//!   positions by `train_len / target_len` before computing RoPE. Simple and
+//!   cheap; slightly lossy at very high extension ratios.
+//!
+//! - **Dynamic NTK scaling** (bloc97 2023): increases the RoPE base
+//!   exponentially so that high-frequency components are scaled less
+//!   aggressively than low-frequency ones.
+//!
+//! - **YaRN** (Peng et al. 2023, arXiv:2309.00071): per-frequency
+//!   interpolation — low-frequency dimensions use linear interpolation,
+//!   high-frequency dimensions are left unscaled, with a smooth ramp between.
+//!   An attention temperature correction factor `√(1/log(target/train))` is
+//!   also computed (callers apply it to attention logit scores).
+//!
+//! # Examples
+//!
+//! ```rust
+//! use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeTable};
+//!
+//! // Standard RoPE table for head_dim=64, 2048 positions
+//! let cfg = RopeScalingConfig::standard(64);
+//! let table = RopeTable::build(&cfg);
+//! assert_eq!(table.cos.len(), 2048 * 32);
+//!
+//! // Apply to a single query vector at position 7
+//! let q = vec![1.0_f32; 64];
+//! let q_rot = table.apply(&q, 7);
+//! assert_eq!(q_rot.len(), 64);
+//! ```
+
+// ============================================================================
+// RopeScalingMode
+// ============================================================================
+
+/// Selects which frequency-scaling strategy to use when building a [`RopeTable`].
+#[derive(Debug, Clone, PartialEq)]
+pub enum RopeScalingMode {
+    /// Standard RoPE — no position scaling.
+    None,
+
+    /// Linear position interpolation: effective position → `position / factor`.
+    ///
+    /// `factor` should be `target_len / train_len` (> 1 when extending context).
+    Linear {
+        /// Scaling factor: positions are divided by this value before computing
+        /// angles, effectively compressing the positional range.
+        factor: f32,
+    },
+
+    /// Dynamic NTK: raises the RoPE base so that high-frequency (low-index)
+    /// dimension pairs are scaled less than low-frequency (high-index) ones.
+    DynamicNtk {
+        /// Original RoPE base (usually 10 000).
+        base: f32,
+        /// The sequence length the model was trained on.
+        train_len: usize,
+    },
+
+    /// YaRN: per-frequency interpolation with attention temperature correction.
+    ///
+    /// - Dimensions with normalised index < `alpha` are left at native
+    ///   frequency (high-frequency, short-range).
+    /// - Dimensions with normalised index > `beta` are linearly interpolated
+    ///   (low-frequency, long-range).
+    /// - In between, a smooth linear ramp blends the two.
+    Yarn {
+        /// Original RoPE base (usually 10 000).
+        base: f32,
+        /// The sequence length the model was trained on.
+        train_len: usize,
+        /// High-frequency threshold (default **1.0**).
+        alpha: f32,
+        /// Low-frequency threshold (default **32.0**).
+        beta: f32,
+    },
+}
+
+// ============================================================================
+// RopeScalingConfig
+// ============================================================================
+
+/// Configuration for building a precomputed [`RopeTable`].
+#[derive(Debug, Clone)]
+pub struct RopeScalingConfig {
+    /// Dimension of each attention head (must be even).
+    pub head_dim: usize,
+    /// Maximum sequence length to pre-compute (number of rows in the table).
+    pub max_position: usize,
+    /// RoPE base frequency (default `10_000.0`).
+    pub base: f32,
+    /// Which scaling strategy to apply.
+    pub mode: RopeScalingMode,
+}
+
+impl RopeScalingConfig {
+    /// Standard (unscaled) RoPE with `base = 10 000` and `max_position = 2048`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::RopeScalingConfig;
+    /// let cfg = RopeScalingConfig::standard(64);
+    /// assert_eq!(cfg.head_dim, 64);
+    /// assert_eq!(cfg.base, 10_000.0);
+    /// ```
+    #[must_use]
+    pub fn standard(head_dim: usize) -> Self {
+        Self {
+            head_dim,
+            max_position: 2048,
+            base: 10_000.0,
+            mode: RopeScalingMode::None,
+        }
+    }
+
+    /// Linear interpolation with the supplied `factor`.
+    ///
+    /// Set `factor = target_len / train_len` to extend to `target_len`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeScalingMode};
+    /// let cfg = RopeScalingConfig::linear(64, 4.0);
+    /// assert!(matches!(cfg.mode, RopeScalingMode::Linear { factor } if factor == 4.0));
+    /// ```
+    #[must_use]
+    pub fn linear(head_dim: usize, factor: f32) -> Self {
+        Self {
+            head_dim,
+            max_position: 2048,
+            base: 10_000.0,
+            mode: RopeScalingMode::Linear { factor },
+        }
+    }
+
+    /// Dynamic NTK scaling.
+    ///
+    /// `train_len` is the sequence length seen during pre-training; `max_pos`
+    /// is the target context length you want the table to cover.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeScalingMode};
+    /// let cfg = RopeScalingConfig::dynamic_ntk(64, 2048, 8192);
+    /// assert!(matches!(cfg.mode, RopeScalingMode::DynamicNtk { .. }));
+    /// ```
+    #[must_use]
+    pub fn dynamic_ntk(head_dim: usize, train_len: usize, max_pos: usize) -> Self {
+        Self {
+            head_dim,
+            max_position: max_pos,
+            base: 10_000.0,
+            mode: RopeScalingMode::DynamicNtk {
+                base: 10_000.0,
+                train_len,
+            },
+        }
+    }
+
+    /// YaRN scaling with default thresholds (`alpha = 1.0`, `beta = 32.0`).
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeScalingMode};
+    /// let cfg = RopeScalingConfig::yarn(64, 2048, 8192);
+    /// assert!(matches!(cfg.mode, RopeScalingMode::Yarn { .. }));
+    /// ```
+    #[must_use]
+    pub fn yarn(head_dim: usize, train_len: usize, max_pos: usize) -> Self {
+        Self {
+            head_dim,
+            max_position: max_pos,
+            base: 10_000.0,
+            mode: RopeScalingMode::Yarn {
+                base: 10_000.0,
+                train_len,
+                alpha: 1.0,
+                beta: 32.0,
+            },
+        }
+    }
+}
+
+// ============================================================================
+// RopeScaler — stateless frequency computations
+// ============================================================================
+
+/// Stateless helper for computing per-dimension-pair inverse frequencies.
+pub struct RopeScaler;
+
+impl RopeScaler {
+    /// Compute the effective inverse frequency for each of the `head_dim / 2`
+    /// rotation pairs, incorporating the chosen scaling strategy.
+    ///
+    /// The returned vector has length `head_dim / 2`. Entry `i` gives the
+    /// angular speed (radians per position unit) for the `i`-th pair of
+    /// dimensions `(2i, 2i+1)`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `config.head_dim` is zero or odd.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeScaler};
+    /// let cfg = RopeScalingConfig::standard(8);
+    /// let inv = RopeScaler::compute_inv_freq(&cfg);
+    /// assert_eq!(inv.len(), 4);
+    /// assert!((inv[0] - 1.0).abs() < 1e-6, "pair 0 freq should be 1.0");
+    /// ```
+    #[must_use]
+    pub fn compute_inv_freq(config: &RopeScalingConfig) -> Vec<f32> {
+        assert!(
+            config.head_dim > 0 && config.head_dim % 2 == 0,
+            "head_dim must be a positive even number, got {}",
+            config.head_dim
+        );
+
+        let half_dim = config.head_dim / 2;
+
+        match &config.mode {
+            // ----------------------------------------------------------------
+            // Standard — unscaled inverse frequencies
+            // ----------------------------------------------------------------
+            RopeScalingMode::None => {
+                (0..half_dim)
+                    .map(|i| {
+                        let exponent = (2 * i) as f32 / config.head_dim as f32;
+                        1.0 / config.base.powf(exponent)
+                    })
+                    .collect()
+            }
+
+            // ----------------------------------------------------------------
+            // Linear — inv_freq is identical to standard; positions are scaled
+            // at table-build time (see RopeTable::build).
+            // ----------------------------------------------------------------
+            RopeScalingMode::Linear { .. } => {
+                (0..half_dim)
+                    .map(|i| {
+                        let exponent = (2 * i) as f32 / config.head_dim as f32;
+                        1.0 / config.base.powf(exponent)
+                    })
+                    .collect()
+            }
+
+            // ----------------------------------------------------------------
+            // Dynamic NTK — raise the effective base by scale^(dim/(dim-2))
+            // ----------------------------------------------------------------
+            RopeScalingMode::DynamicNtk { base, train_len } => {
+                let s = config.max_position as f32 / *train_len as f32;
+                let dim_ratio = config.head_dim as f32 / (config.head_dim as f32 - 2.0);
+                let base_new = base * s.powf(dim_ratio);
+                (0..half_dim)
+                    .map(|i| {
+                        let exponent = (2 * i) as f32 / config.head_dim as f32;
+                        1.0 / base_new.powf(exponent)
+                    })
+                    .collect()
+            }
+
+            // ----------------------------------------------------------------
+            // YaRN — per-dimension blend of unscaled and linearly scaled freq
+            // ----------------------------------------------------------------
+            RopeScalingMode::Yarn {
+                base,
+                train_len,
+                alpha,
+                beta,
+            } => {
+                let scale = config.max_position as f32 / *train_len as f32;
+                (0..half_dim)
+                    .map(|i| {
+                        // Standard inverse frequency for this pair
+                        let exponent = (2 * i) as f32 / config.head_dim as f32;
+                        let freq_std = 1.0 / base.powf(exponent);
+
+                        // Normalised dimension index in [0, half_dim)
+                        let d = i as f32;
+                        let ramp = Self::yarn_ramp(d, *alpha, *beta);
+
+                        // ramp=0 → unscaled (high-freq), ramp=1 → interpolated (low-freq)
+                        // Effective inv_freq: blend denominator
+                        freq_std / (1.0 - ramp + ramp * scale)
+                    })
+                    .collect()
+            }
+        }
+    }
+
+    /// YaRN ramp function: returns how much linear interpolation to apply to
+    /// dimension pair `d`.
+    ///
+    /// - Returns `0.0` for high-frequency dimensions (`d ≤ alpha`).
+    /// - Returns `1.0` for low-frequency dimensions (`d ≥ beta`).
+    /// - Linearly interpolates between `alpha` and `beta`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::RopeScaler;
+    /// assert_eq!(RopeScaler::yarn_ramp(0.0, 1.0, 32.0), 0.0);
+    /// assert_eq!(RopeScaler::yarn_ramp(32.0, 1.0, 32.0), 1.0);
+    /// assert!((RopeScaler::yarn_ramp(16.5, 1.0, 32.0) - 0.5).abs() < 1e-5);
+    /// ```
+    #[must_use]
+    pub fn yarn_ramp(d: f32, alpha: f32, beta: f32) -> f32 {
+        if (beta - alpha).abs() < f32::EPSILON {
+            // Degenerate interval: treat everything as low-freq
+            return if d >= beta { 1.0 } else { 0.0 };
+        }
+        let r = (d - alpha) / (beta - alpha);
+        r.clamp(0.0, 1.0)
+    }
+
+    /// Attention temperature correction factor for YaRN.
+    ///
+    /// Returns `√(1 / log(target_len / train_len))`. Callers multiply raw
+    /// attention logit scores by this value before softmax.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `target_len ≤ train_len` (no extension, temperature undefined)
+    /// or if either length is zero.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::RopeScaler;
+    /// let t = RopeScaler::yarn_temperature(8192, 2048);
+    /// // log(4) ≈ 1.386; sqrt(1/1.386) ≈ 0.849
+    /// assert!((t - (1.0_f32 / (4.0_f32.ln())).sqrt()).abs() < 1e-5);
+    /// ```
+    #[must_use]
+    pub fn yarn_temperature(target_len: usize, train_len: usize) -> f32 {
+        assert!(
+            target_len > train_len && train_len > 0,
+            "target_len ({target_len}) must be strictly greater than train_len ({train_len})"
+        );
+        let ratio = target_len as f32 / train_len as f32;
+        (1.0 / ratio.ln()).sqrt()
+    }
+}
+
+// ============================================================================
+// RopeTable — precomputed cosine/sine lookup tables
+// ============================================================================
+
+/// Precomputed RoPE cosine and sine tables for efficient inference.
+///
+/// Layout: both `cos` and `sin` are flat vectors of length
+/// `max_position × (head_dim / 2)`.  Entry `[pos * half_dim + i]` holds the
+/// cosine (or sine) of the angle for position `pos` and dimension pair `i`.
+///
+/// # Examples
+///
+/// ```rust
+/// use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeTable};
+/// let cfg = RopeScalingConfig::standard(16);
+/// let table = RopeTable::build(&cfg);
+/// // Position 0 must be an identity rotation (cos=1, sin=0)
+/// assert!((table.cos[0] - 1.0).abs() < 1e-6);
+/// assert!((table.sin[0]).abs() < 1e-6);
+/// ```
+#[derive(Debug, Clone)]
+pub struct RopeTable {
+    /// Cosine values, shape `[max_position × head_dim/2]`.
+    pub cos: Vec<f32>,
+    /// Sine values, shape `[max_position × head_dim/2]`.
+    pub sin: Vec<f32>,
+    /// Number of positions in the table.
+    pub max_position: usize,
+    /// Attention-head dimension (must be even).
+    pub head_dim: usize,
+}
+
+impl RopeTable {
+    /// Precompute the RoPE table from `config`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `config.head_dim` is zero or odd.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeTable};
+    /// let cfg = RopeScalingConfig::standard(64);
+    /// let table = RopeTable::build(&cfg);
+    /// assert_eq!(table.cos.len(), 2048 * 32);
+    /// ```
+    #[must_use]
+    pub fn build(config: &RopeScalingConfig) -> Self {
+        let inv_freq = RopeScaler::compute_inv_freq(config);
+        let half_dim = config.head_dim / 2;
+        let max_pos = config.max_position;
+
+        let mut cos = vec![0.0_f32; max_pos * half_dim];
+        let mut sin = vec![0.0_f32; max_pos * half_dim];
+
+        for pos in 0..max_pos {
+            // For Linear mode, compress positions by the scale factor so the
+            // model sees effective positions in [0, train_len) even when
+            // absolute positions exceed train_len.
+            let pos_f = match &config.mode {
+                RopeScalingMode::Linear { factor } => pos as f32 / factor,
+                _ => pos as f32,
+            };
+
+            for i in 0..half_dim {
+                let angle = pos_f * inv_freq[i];
+                cos[pos * half_dim + i] = angle.cos();
+                sin[pos * half_dim + i] = angle.sin();
+            }
+        }
+
+        Self {
+            cos,
+            sin,
+            max_position: max_pos,
+            head_dim: config.head_dim,
+        }
+    }
+
+    /// Apply RoPE to a single head vector `x` of length `head_dim` at the
+    /// given absolute `position`.
+    ///
+    /// The standard complex rotation for dimension pair `i` is:
+    ///
+    /// ```text
+    /// x_out[2i]   = x[2i]   * cos[θ_i] - x[2i+1] * sin[θ_i]
+    /// x_out[2i+1] = x[2i]   * sin[θ_i] + x[2i+1] * cos[θ_i]
+    /// ```
+    ///
+    /// This preserves the L2 norm of `x`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `x.len() != self.head_dim` or `position >= self.max_position`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeTable};
+    /// let table = RopeTable::build(&RopeScalingConfig::standard(8));
+    /// let x = vec![1.0_f32, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0];
+    /// // Position 0: cos=1, sin=0 → identity
+    /// let out = table.apply(&x, 0);
+    /// for (a, b) in x.iter().zip(out.iter()) {
+    ///     assert!((a - b).abs() < 1e-6);
+    /// }
+    /// ```
+    #[must_use]
+    pub fn apply(&self, x: &[f32], position: usize) -> Vec<f32> {
+        assert_eq!(
+            x.len(),
+            self.head_dim,
+            "input length {} does not match head_dim {}",
+            x.len(),
+            self.head_dim
+        );
+        assert!(
+            position < self.max_position,
+            "position {position} >= max_position {}",
+            self.max_position
+        );
+
+        let half_dim = self.head_dim / 2;
+        let mut out = vec![0.0_f32; self.head_dim];
+
+        for i in 0..half_dim {
+            let c = self.cos[position * half_dim + i];
+            let s = self.sin[position * half_dim + i];
+            let x0 = x[2 * i];
+            let x1 = x[2 * i + 1];
+            out[2 * i] = x0 * c - x1 * s;
+            out[2 * i + 1] = x0 * s + x1 * c;
+        }
+
+        out
+    }
+
+    /// Apply RoPE to all tokens in a sequence.
+    ///
+    /// `x` must be in **(seq_len, num_heads, head_dim)** layout (row-major).
+    /// Returns a new allocation with the same shape.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `x.len() != seq_len * num_heads * head_dim`, or if
+    /// `seq_len > self.max_position`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_transformers::layers::rope_scaling::{RopeScalingConfig, RopeTable};
+    /// let table = RopeTable::build(&RopeScalingConfig::standard(8));
+    /// let seq_len = 4;
+    /// let num_heads = 2;
+    /// let x = vec![1.0_f32; seq_len * num_heads * 8];
+    /// let out = table.apply_to_sequence(&x, seq_len, num_heads);
+    /// assert_eq!(out.len(), x.len());
+    /// ```
+    #[must_use]
+    pub fn apply_to_sequence(&self, x: &[f32], seq_len: usize, num_heads: usize) -> Vec<f32> {
+        assert_eq!(
+            x.len(),
+            seq_len * num_heads * self.head_dim,
+            "input length {} != seq_len({}) * num_heads({}) * head_dim({})",
+            x.len(),
+            seq_len,
+            num_heads,
+            self.head_dim
+        );
+        assert!(
+            seq_len <= self.max_position,
+            "seq_len {seq_len} exceeds max_position {}",
+            self.max_position
+        );
+
+        let mut out = vec![0.0_f32; x.len()];
+        let head_stride = self.head_dim;
+        let seq_stride = num_heads * self.head_dim;
+
+        for pos in 0..seq_len {
+            for h in 0..num_heads {
+                let src_offset = pos * seq_stride + h * head_stride;
+                let head_slice = &x[src_offset..src_offset + self.head_dim];
+                let rotated = self.apply(head_slice, pos);
+                out[src_offset..src_offset + self.head_dim].copy_from_slice(&rotated);
+            }
+        }
+
+        out
+    }
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // -----------------------------------------------------------------------
+    // Helper
+    // -----------------------------------------------------------------------
+
+    fn l2_norm(v: &[f32]) -> f32 {
+        v.iter().map(|x| x * x).sum::<f32>().sqrt()
+    }
+
+    // -----------------------------------------------------------------------
+    // Inverse frequency tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_standard_inv_freq_first() {
+        // inv_freq[0] = 1 / base^0 = 1.0
+        let cfg = RopeScalingConfig::standard(16);
+        let inv = RopeScaler::compute_inv_freq(&cfg);
+        assert!(
+            (inv[0] - 1.0).abs() < 1e-6,
+            "inv_freq[0] should be 1.0, got {}",
+            inv[0]
+        );
+    }
+
+    #[test]
+    fn test_standard_inv_freq_decreasing() {
+        // Higher-index pairs have smaller inverse frequency (slower rotation)
+        let cfg = RopeScalingConfig::standard(32);
+        let inv = RopeScaler::compute_inv_freq(&cfg);
+        for i in 1..inv.len() {
+            assert!(
+                inv[i] < inv[i - 1],
+                "inv_freq should be strictly decreasing: inv[{}]={} >= inv[{}]={}",
+                i,
+                inv[i],
+                i - 1,
+                inv[i - 1]
+            );
+        }
+    }
+
+    #[test]
+    fn test_linear_inv_freq_same_as_standard() {
+        // Linear mode: inv_freq is identical to standard — positions are scaled
+        // at table-build time, not by changing frequencies.
+        let cfg_std = RopeScalingConfig::standard(32);
+        let cfg_lin = RopeScalingConfig {
+            mode: RopeScalingMode::Linear { factor: 4.0 },
+            ..cfg_std.clone()
+        };
+        let inv_std = RopeScaler::compute_inv_freq(&cfg_std);
+        let inv_lin = RopeScaler::compute_inv_freq(&cfg_lin);
+        for (a, b) in inv_std.iter().zip(inv_lin.iter()) {
+            assert!(
+                (a - b).abs() < 1e-7,
+                "Linear inv_freq {b} != standard {a}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_dynamic_ntk_larger_base() {
+        // NTK uses a scaled-up base → smaller inv_freq values than standard
+        let cfg_std = RopeScalingConfig::standard(32);
+        let cfg_ntk = RopeScalingConfig::dynamic_ntk(32, 2048, 8192);
+        let inv_std = RopeScaler::compute_inv_freq(&cfg_std);
+        let inv_ntk = RopeScaler::compute_inv_freq(&cfg_ntk);
+        // pair 0 is inv_freq = 1/(base^0) = 1 for both — compare pair 1+
+        for i in 1..inv_std.len() {
+            assert!(
+                inv_ntk[i] < inv_std[i],
+                "NTK inv_freq[{}]={} should be < standard {}",
+                i,
+                inv_ntk[i],
+                inv_std[i]
+            );
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // YaRN ramp tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_yarn_ramp_zero_at_high_freq() {
+        // d=0 (< alpha=1) → ramp == 0 (no interpolation, keep native freq)
+        let r = RopeScaler::yarn_ramp(0.0, 1.0, 32.0);
+        assert_eq!(r, 0.0, "ramp at d=0 should be 0.0, got {r}");
+    }
+
+    #[test]
+    fn test_yarn_ramp_one_at_low_freq() {
+        // d=32 (== beta=32) → ramp == 1 (full linear interpolation)
+        let r = RopeScaler::yarn_ramp(32.0, 1.0, 32.0);
+        assert_eq!(r, 1.0, "ramp at d=32 should be 1.0, got {r}");
+    }
+
+    #[test]
+    fn test_yarn_ramp_midpoint() {
+        // d midway between alpha=1 and beta=32 → ramp ≈ 0.5
+        let mid = (1.0 + 32.0) / 2.0;
+        let r = RopeScaler::yarn_ramp(mid, 1.0, 32.0);
+        assert!(
+            (r - 0.5).abs() < 1e-5,
+            "ramp at midpoint should be ~0.5, got {r}"
+        );
+    }
+
+    #[test]
+    fn test_yarn_ramp_clamps_below_zero() {
+        // d well below alpha → must not return negative values
+        let r = RopeScaler::yarn_ramp(-100.0, 1.0, 32.0);
+        assert_eq!(r, 0.0, "ramp below alpha must clamp to 0.0, got {r}");
+    }
+
+    #[test]
+    fn test_yarn_ramp_clamps_above_one() {
+        // d well above beta → must not exceed 1.0
+        let r = RopeScaler::yarn_ramp(1000.0, 1.0, 32.0);
+        assert_eq!(r, 1.0, "ramp above beta must clamp to 1.0, got {r}");
+    }
+
+    // -----------------------------------------------------------------------
+    // YaRN temperature tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_yarn_temperature_known_value() {
+        // target=8192, train=2048, ratio=4 → sqrt(1/ln(4))
+        let t = RopeScaler::yarn_temperature(8192, 2048);
+        let expected = (1.0_f32 / 4.0_f32.ln()).sqrt();
+        assert!(
+            (t - expected).abs() < 1e-5,
+            "temperature {t} != expected {expected}"
+        );
+    }
+
+    #[test]
+    fn test_yarn_temperature_decreases_with_scale() {
+        // Larger extension ratio → log grows → temperature shrinks
+        let t_small = RopeScaler::yarn_temperature(4096, 2048); // ratio 2
+        let t_large = RopeScaler::yarn_temperature(8192, 2048); // ratio 4
+        assert!(
+            t_large < t_small,
+            "temperature should decrease with larger scale: {t_large} >= {t_small}"
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Table shape tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_table_build_shape() {
+        let cfg = RopeScalingConfig {
+            head_dim: 32,
+            max_position: 128,
+            base: 10_000.0,
+            mode: RopeScalingMode::None,
+        };
+        let table = RopeTable::build(&cfg);
+        assert_eq!(table.cos.len(), 128 * 16, "cos shape mismatch");
+        assert_eq!(table.sin.len(), 128 * 16, "sin shape mismatch");
+        assert_eq!(table.max_position, 128);
+        assert_eq!(table.head_dim, 32);
+    }
+
+    #[test]
+    fn test_table_cos_sin_unit_circle() {
+        // For every entry: cos² + sin² must equal 1
+        let cfg = RopeScalingConfig::standard(16);
+        let table = RopeTable::build(&cfg);
+        for (c, s) in table.cos.iter().zip(table.sin.iter()) {
+            let r = c * c + s * s;
+            assert!(
+                (r - 1.0).abs() < 1e-5,
+                "cos²+sin²={r} at some entry (expected 1.0)"
+            );
+        }
+    }
+
+    #[test]
+    fn test_table_position_zero() {
+        // At position 0, angle = 0 → cos=1, sin=0 for every dimension pair
+        let cfg = RopeScalingConfig::standard(16);
+        let table = RopeTable::build(&cfg);
+        let half_dim = 8;
+        for i in 0..half_dim {
+            let c = table.cos[i];
+            let s = table.sin[i];
+            assert!(
+                (c - 1.0).abs() < 1e-6,
+                "cos at pos=0, pair {i} should be 1.0, got {c}"
+            );
+            assert!(
+                s.abs() < 1e-6,
+                "sin at pos=0, pair {i} should be 0.0, got {s}"
+            );
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Apply (single vector) tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_apply_rotation_shape() {
+        let table = RopeTable::build(&RopeScalingConfig::standard(32));
+        let x = vec![0.5_f32; 32];
+        let out = table.apply(&x, 0);
+        assert_eq!(out.len(), 32, "output length should equal head_dim");
+    }
+
+    #[test]
+    fn test_apply_position_zero_identity() {
+        // At position 0 (cos=1, sin=0) the rotation is the identity
+        let table = RopeTable::build(&RopeScalingConfig::standard(16));
+        let x: Vec<f32> = (0..16).map(|i| i as f32 + 0.1).collect();
+        let out = table.apply(&x, 0);
+        for (a, b) in x.iter().zip(out.iter()) {
+            assert!(
+                (a - b).abs() < 1e-6,
+                "position 0 should be identity: in={a}, out={b}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_apply_rotation_preserves_norm() {
+        // RoPE is a unitary rotation → it must preserve the L2 norm
+        let table = RopeTable::build(&RopeScalingConfig::standard(32));
+        let x: Vec<f32> = (0..32).map(|i| (i as f32 + 1.0).recip()).collect();
+        for pos in [0_usize, 1, 7, 31, 100] {
+            let out = table.apply(&x, pos);
+            let norm_in = l2_norm(&x);
+            let norm_out = l2_norm(&out);
+            assert!(
+                (norm_in - norm_out).abs() < 1e-4,
+                "norm not preserved at pos {pos}: in={norm_in:.6} out={norm_out:.6}"
+            );
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Apply-to-sequence tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_apply_sequence_shape() {
+        let table = RopeTable::build(&RopeScalingConfig::standard(16));
+        let seq_len = 8;
+        let num_heads = 4;
+        let x = vec![1.0_f32; seq_len * num_heads * 16];
+        let out = table.apply_to_sequence(&x, seq_len, num_heads);
+        assert_eq!(out.len(), x.len(), "sequence output must keep same shape");
+    }
+
+    #[test]
+    fn test_apply_sequence_correct_per_position() {
+        // Each position-row in the sequence output must match apply(·, pos)
+        let cfg = RopeScalingConfig::standard(8);
+        let table = RopeTable::build(&cfg);
+        let seq_len = 5;
+        let num_heads = 2;
+        let x: Vec<f32> = (0..seq_len * num_heads * 8)
+            .map(|i| i as f32 * 0.1 + 0.5)
+            .collect();
+
+        let out_seq = table.apply_to_sequence(&x, seq_len, num_heads);
+
+        // Check position 1 (arbitrary)
+        let pos = 1;
+        for h in 0..num_heads {
+            let src = pos * num_heads * 8 + h * 8;
+            let expected = table.apply(&x[src..src + 8], pos);
+            let actual = &out_seq[src..src + 8];
+            for (e, a) in expected.iter().zip(actual.iter()) {
+                assert!(
+                    (e - a).abs() < 1e-6,
+                    "sequence[pos={pos}, head={h}] mismatch: expected {e}, got {a}"
+                );
+            }
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Scaling-specific correctness tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_linear_scales_positions() {
+        // Linear factor=2: position 10 in the scaled table should see angles
+        // identical to position 5 in a standard table (10/2 = 5).
+        let cfg_std = RopeScalingConfig {
+            max_position: 512,
+            ..RopeScalingConfig::standard(16)
+        };
+        let cfg_lin = RopeScalingConfig {
+            max_position: 512,
+            mode: RopeScalingMode::Linear { factor: 2.0 },
+            ..RopeScalingConfig::standard(16)
+        };
+        let table_std = RopeTable::build(&cfg_std);
+        let table_lin = RopeTable::build(&cfg_lin);
+
+        let half = 8;
+        // Standard at pos 5 vs Linear at pos 10
+        for i in 0..half {
+            let c_std = table_std.cos[5 * half + i];
+            let c_lin = table_lin.cos[10 * half + i];
+            assert!(
+                (c_std - c_lin).abs() < 1e-5,
+                "cos: std[5][{i}]={c_std} != lin[10][{i}]={c_lin}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_standard_vs_linear_position_compression() {
+        // Alias of test_linear_scales_positions with sine channel
+        let cfg_std = RopeScalingConfig {
+            max_position: 256,
+            ..RopeScalingConfig::standard(16)
+        };
+        let cfg_lin = RopeScalingConfig {
+            max_position: 256,
+            mode: RopeScalingMode::Linear { factor: 2.0 },
+            ..RopeScalingConfig::standard(16)
+        };
+        let t_std = RopeTable::build(&cfg_std);
+        let t_lin = RopeTable::build(&cfg_lin);
+        let half = 8;
+        for i in 0..half {
+            let s_std = t_std.sin[5 * half + i];
+            let s_lin = t_lin.sin[10 * half + i];
+            assert!(
+                (s_std - s_lin).abs() < 1e-5,
+                "sin mismatch: std[5][{i}]={s_std} != lin[10][{i}]={s_lin}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_dynamic_ntk_high_freq_small() {
+        // NTK raises the effective base → inv_freq values are strictly smaller
+        // than standard for all non-trivial pairs (pairs 1..).
+        // This is already proven by comparing inv_freq directly (avoiding the
+        // ambiguity of cosine wrapping at large position × frequency products).
+        let cfg_std = RopeScalingConfig {
+            max_position: 8192,
+            ..RopeScalingConfig::standard(32)
+        };
+        let cfg_ntk = RopeScalingConfig {
+            max_position: 8192,
+            ..RopeScalingConfig::dynamic_ntk(32, 2048, 8192)
+        };
+        let inv_std = RopeScaler::compute_inv_freq(&cfg_std);
+        let inv_ntk = RopeScaler::compute_inv_freq(&cfg_ntk);
+
+        // For every pair >= 1 the NTK frequency must be strictly below standard.
+        for i in 1..inv_std.len() {
+            assert!(
+                inv_ntk[i] < inv_std[i],
+                "NTK inv_freq[{i}]={} should be < standard {} (larger base → smaller freq)",
+                inv_ntk[i],
+                inv_std[i]
+            );
+        }
+
+        // And at a small position (pos=1) where cos has not wrapped, the NTK
+        // angle is indeed smaller (cos closer to 1) for the highest-frequency
+        // pair that still differs meaningfully.
+        let half = 16_usize;
+        let pos = 1_usize;
+        // inv_freq[0] = 1.0 for both (base^0 = 1), so angle at pos=1 is just
+        // 1.0 rad for both — pick pair 1 instead.
+        let angle_std = inv_std[1] * pos as f32;
+        let angle_ntk = inv_ntk[1] * pos as f32;
+        // Both angles are < π/2 for pair 1 at pos=1; cos should be larger for NTK.
+        assert!(
+            angle_ntk < angle_std,
+            "NTK angle at pos={pos}, pair=1 should be smaller: ntk={angle_ntk}, std={angle_std}"
+        );
+        drop(half); // suppress unused variable warning
+    }
+
+    #[test]
+    fn test_yarn_low_freq_interpolated() {
+        // YaRN low-frequency dimensions (high dim index) are interpolated:
+        // they should behave similarly to Linear-scaled positions.
+        // For the LAST pair (index half_dim-1), ramp ≈ 1 → fully interpolated.
+        let head_dim = 64usize;
+        let half = head_dim / 2;
+        let train_len = 2048usize;
+        let target = 4096usize;
+        let scale = target as f32 / train_len as f32; // = 2.0
+
+        let cfg_yarn = RopeScalingConfig {
+            head_dim,
+            max_position: target,
+            base: 10_000.0,
+            mode: RopeScalingMode::Yarn {
+                base: 10_000.0,
+                train_len,
+                alpha: 1.0,
+                beta: 32.0,
+            },
+        };
+        // Equivalent linear config (factor = scale)
+        let cfg_lin = RopeScalingConfig {
+            head_dim,
+            max_position: target,
+            base: 10_000.0,
+            mode: RopeScalingMode::Linear { factor: scale },
+        };
+
+        let t_yarn = RopeTable::build(&cfg_yarn);
+        let t_lin = RopeTable::build(&cfg_lin);
+
+        // At the last pair (lowest frequency), YaRN ≈ linear interpolation.
+        // Compare cos at position 100 for the last dim pair.
+        let pos = 100;
+        let i = half - 1; // last pair
+        let c_yarn = t_yarn.cos[pos * half + i];
+        let c_lin = t_lin.cos[pos * half + i];
+        // Should be close (tolerance allows for ramp not being exactly 1.0
+        // at the boundary since beta=32 and half_dim-1=31)
+        assert!(
+            (c_yarn - c_lin).abs() < 0.01,
+            "YaRN last dim should approximate linear: yarn={c_yarn}, lin={c_lin}"
+        );
+    }
+}

crates/training/rtx-transformers/src/training/dpo.rs

@@ -0,0 +1,663 @@
+//! Direct Preference Optimization (DPO) loss.
+//!
+//! Implements DPO (Rafailov et al. 2023, arXiv:2305.18290) and IPO (Azar et al. 2023,
+//! arXiv:2310.12036). Both methods train language models on preference data (chosen vs
+//! rejected responses) without a separate reward model.
+//!
+//! # DPO loss (per sample)
+//!
+//! ```text
+//! log_ratio_chosen   = log π_θ(y_w|x) - log π_ref(y_w|x)
+//! log_ratio_rejected = log π_θ(y_l|x) - log π_ref(y_l|x)
+//! reward_margin      = β * (log_ratio_chosen - log_ratio_rejected)
+//! loss               = -log σ(reward_margin)
+//! ```
+//!
+//! # IPO loss (per sample)
+//!
+//! ```text
+//! loss = (log_ratio_chosen - log_ratio_rejected - 1/(2β))²
+//! ```
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::training::dpo::{DpoLoss, DpoConfig, DpoVariant};
+//!
+//! let loss_fn = DpoLoss::with_beta(0.1);
+//!
+//! // chosen response has higher log-probability under policy
+//! let result = loss_fn.compute(-2.0, -5.0, -3.0, -4.0);
+//! assert!(result.loss < 0.693, "loss should be below log(2) when chosen is preferred");
+//! assert!(result.chosen_preferred);
+//! ```
+
+/// Which variant of the preference optimization loss to use.
+#[derive(Debug, Clone, PartialEq)]
+pub enum DpoVariant {
+    /// Standard DPO — log-sigmoid loss (Rafailov et al. 2023).
+    Dpo,
+    /// IPO — squared loss that avoids the log-sigmoid collapse problem (Azar et al. 2023).
+    Ipo,
+}
+
+/// Configuration for the DPO/IPO loss function.
+#[derive(Debug, Clone)]
+pub struct DpoConfig {
+    /// KL regularisation temperature. Controls how far the policy may drift from the
+    /// reference model. Typical values: 0.01–0.5. Default: 0.1.
+    pub beta: f32,
+    /// Which loss variant to use. Default: [`DpoVariant::Dpo`].
+    pub variant: DpoVariant,
+    /// Label smoothing coefficient ε for Robust DPO (Mitchell et al. 2023).
+    ///
+    /// When `label_smoothing > 0.0` the DPO loss becomes:
+    /// ```text
+    /// (1 - ε) * (-log σ(margin)) + ε * (-log σ(-margin))
+    /// ```
+    /// Has no effect when `variant == DpoVariant::Ipo`. Default: 0.0.
+    pub label_smoothing: f32,
+}
+
+impl Default for DpoConfig {
+    fn default() -> Self {
+        Self {
+            beta: 0.1,
+            variant: DpoVariant::Dpo,
+            label_smoothing: 0.0,
+        }
+    }
+}
+
+/// Per-sample result from the DPO/IPO loss computation.
+#[derive(Debug, Clone)]
+pub struct DpoLossResult {
+    /// Scalar loss value for this preference pair.
+    pub loss: f32,
+    /// Implicit reward for the chosen response: β * (log π - log π_ref).
+    pub reward_chosen: f32,
+    /// Implicit reward for the rejected response: β * (log π - log π_ref).
+    pub reward_rejected: f32,
+    /// `reward_chosen - reward_rejected`.
+    pub reward_margin: f32,
+    /// `true` when `reward_chosen > reward_rejected` (model prefers the correct response).
+    pub chosen_preferred: bool,
+}
+
+/// Running accumulator for batch-level DPO metrics.
+///
+/// All sums are kept in `f64` to avoid catastrophic cancellation over large batches.
+#[derive(Debug, Clone, Default)]
+pub struct DpoAccumulator {
+    total_loss: f64,
+    total_reward_chosen: f64,
+    total_reward_rejected: f64,
+    total_reward_margin: f64,
+    n_preferred: usize,
+    n_samples: usize,
+}
+
+impl DpoAccumulator {
+    /// Create an empty accumulator.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Incorporate one [`DpoLossResult`] into the running totals.
+    pub fn add(&mut self, result: &DpoLossResult) {
+        self.total_loss += result.loss as f64;
+        self.total_reward_chosen += result.reward_chosen as f64;
+        self.total_reward_rejected += result.reward_rejected as f64;
+        self.total_reward_margin += result.reward_margin as f64;
+        if result.chosen_preferred {
+            self.n_preferred += 1;
+        }
+        self.n_samples += 1;
+    }
+
+    /// Mean loss across all accumulated samples. Returns 0.0 if no samples have been added.
+    pub fn avg_loss(&self) -> f32 {
+        if self.n_samples == 0 {
+            return 0.0;
+        }
+        (self.total_loss / self.n_samples as f64) as f32
+    }
+
+    /// Mean implicit reward for chosen responses.
+    pub fn avg_reward_chosen(&self) -> f32 {
+        if self.n_samples == 0 {
+            return 0.0;
+        }
+        (self.total_reward_chosen / self.n_samples as f64) as f32
+    }
+
+    /// Mean implicit reward for rejected responses.
+    pub fn avg_reward_rejected(&self) -> f32 {
+        if self.n_samples == 0 {
+            return 0.0;
+        }
+        (self.total_reward_rejected / self.n_samples as f64) as f32
+    }
+
+    /// Mean reward margin (chosen − rejected).
+    pub fn avg_reward_margin(&self) -> f32 {
+        if self.n_samples == 0 {
+            return 0.0;
+        }
+        (self.total_reward_margin / self.n_samples as f64) as f32
+    }
+
+    /// Fraction of samples where the model correctly preferred the chosen response.
+    /// Returns 0.0 if no samples have been added.
+    pub fn preference_accuracy(&self) -> f32 {
+        if self.n_samples == 0 {
+            return 0.0;
+        }
+        self.n_preferred as f32 / self.n_samples as f32
+    }
+
+    /// Total number of samples accumulated so far.
+    pub fn n_samples(&self) -> usize {
+        self.n_samples
+    }
+
+    /// Reset all accumulators to zero.
+    pub fn reset(&mut self) {
+        *self = Self::default();
+    }
+}
+
+/// DPO/IPO loss function.
+///
+/// # Example
+///
+/// ```rust
+/// use rtx_transformers::training::dpo::{DpoLoss, DpoConfig, DpoVariant};
+///
+/// // Standard DPO
+/// let dpo = DpoLoss::with_beta(0.1);
+/// let result = dpo.compute(-1.0, -4.0, -2.0, -3.0);
+/// println!("loss={:.4}  margin={:.4}  acc={}", result.loss, result.reward_margin, result.chosen_preferred);
+///
+/// // IPO variant
+/// let ipo = DpoLoss::ipo(0.1);
+/// let result = ipo.compute(-1.0, -4.0, -2.0, -3.0);
+/// println!("ipo loss={:.4}", result.loss);
+/// ```
+pub struct DpoLoss {
+    config: DpoConfig,
+}
+
+impl DpoLoss {
+    /// Create a [`DpoLoss`] from an explicit [`DpoConfig`].
+    pub fn new(config: DpoConfig) -> Self {
+        Self { config }
+    }
+
+    /// Create a [`DpoLoss`] with default configuration (DPO, β=0.1, no smoothing).
+    pub fn default() -> Self {
+        Self::new(DpoConfig::default())
+    }
+
+    /// Create a standard DPO loss with the given β.
+    pub fn with_beta(beta: f32) -> Self {
+        Self::new(DpoConfig {
+            beta,
+            ..DpoConfig::default()
+        })
+    }
+
+    /// Create an IPO loss with the given β.
+    pub fn ipo(beta: f32) -> Self {
+        Self::new(DpoConfig {
+            beta,
+            variant: DpoVariant::Ipo,
+            label_smoothing: 0.0,
+        })
+    }
+
+    /// Compute the DPO/IPO loss for one preference pair.
+    ///
+    /// All log-probability arguments should be **sum** log-probs over the response tokens
+    /// (i.e., negative values whose magnitude grows with sequence length).
+    ///
+    /// # Arguments
+    ///
+    /// * `policy_logp_chosen`   — log π_θ(y_w | x)
+    /// * `policy_logp_rejected` — log π_θ(y_l | x)
+    /// * `ref_logp_chosen`      — log π_ref(y_w | x)
+    /// * `ref_logp_rejected`    — log π_ref(y_l | x)
+    pub fn compute(
+        &self,
+        policy_logp_chosen: f32,
+        policy_logp_rejected: f32,
+        ref_logp_chosen: f32,
+        ref_logp_rejected: f32,
+    ) -> DpoLossResult {
+        let beta = self.config.beta;
+
+        let log_ratio_chosen = policy_logp_chosen - ref_logp_chosen;
+        let log_ratio_rejected = policy_logp_rejected - ref_logp_rejected;
+
+        let reward_chosen = beta * log_ratio_chosen;
+        let reward_rejected = beta * log_ratio_rejected;
+        let reward_margin = reward_chosen - reward_rejected;
+
+        let loss = match self.config.variant {
+            DpoVariant::Dpo => {
+                let eps = self.config.label_smoothing;
+                let base_loss = -Self::log_sigmoid(reward_margin);
+                if eps > 0.0 {
+                    // Robust DPO: mix in the reversed direction
+                    (1.0 - eps) * base_loss + eps * (-Self::log_sigmoid(-reward_margin))
+                } else {
+                    base_loss
+                }
+            }
+            DpoVariant::Ipo => {
+                let target = 1.0 / (2.0 * beta);
+                let diff = log_ratio_chosen - log_ratio_rejected - target;
+                diff * diff
+            }
+        };
+
+        DpoLossResult {
+            loss,
+            reward_chosen,
+            reward_rejected,
+            reward_margin,
+            chosen_preferred: reward_chosen > reward_rejected,
+        }
+    }
+
+    /// Compute the loss for a batch of preference pairs.
+    ///
+    /// All four slices must have the same length. Returns the mean loss across the batch
+    /// together with a filled [`DpoAccumulator`].
+    ///
+    /// # Panics
+    ///
+    /// Panics if the slice lengths differ.
+    pub fn compute_batch(
+        &self,
+        policy_logps_chosen: &[f32],
+        policy_logps_rejected: &[f32],
+        ref_logps_chosen: &[f32],
+        ref_logps_rejected: &[f32],
+    ) -> (f32, DpoAccumulator) {
+        assert_eq!(
+            policy_logps_chosen.len(),
+            policy_logps_rejected.len(),
+            "batch slice length mismatch: policy chosen vs rejected"
+        );
+        assert_eq!(
+            policy_logps_chosen.len(),
+            ref_logps_chosen.len(),
+            "batch slice length mismatch: policy chosen vs ref chosen"
+        );
+        assert_eq!(
+            policy_logps_chosen.len(),
+            ref_logps_rejected.len(),
+            "batch slice length mismatch: policy chosen vs ref rejected"
+        );
+
+        let mut acc = DpoAccumulator::new();
+        for i in 0..policy_logps_chosen.len() {
+            let result = self.compute(
+                policy_logps_chosen[i],
+                policy_logps_rejected[i],
+                ref_logps_chosen[i],
+                ref_logps_rejected[i],
+            );
+            acc.add(&result);
+        }
+        let avg = acc.avg_loss();
+        (avg, acc)
+    }
+
+    /// Numerically stable log-sigmoid: log(σ(x)) = −log(1 + exp(−x)).
+    ///
+    /// Uses the identity:
+    /// ```text
+    /// x ≥ 0:  log σ(x) = −log(1 + exp(−x))
+    /// x < 0:  log σ(x) = x − log(1 + exp(x))
+    /// ```
+    /// This avoids overflow for large |x| while preserving full f32 precision.
+    pub fn log_sigmoid(x: f32) -> f32 {
+        if x >= 0.0 {
+            -(1.0_f32 + (-x).exp()).ln()
+        } else {
+            x - (1.0_f32 + x.exp()).ln()
+        }
+    }
+
+    /// Compute the implicit reward β * (log π_θ(y|x) − log π_ref(y|x)).
+    pub fn implicit_reward(beta: f32, policy_logp: f32, ref_logp: f32) -> f32 {
+        beta * (policy_logp - ref_logp)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const EPS: f32 = 1e-5;
+
+    fn approx_eq(a: f32, b: f32, tol: f32) -> bool {
+        (a - b).abs() <= tol
+    }
+
+    // ------------------------------------------------------------------
+    // Configuration constructors
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_default_config() {
+        let cfg = DpoConfig::default();
+        assert!(approx_eq(cfg.beta, 0.1, EPS));
+        assert_eq!(cfg.variant, DpoVariant::Dpo);
+        assert!(approx_eq(cfg.label_smoothing, 0.0, EPS));
+    }
+
+    #[test]
+    fn test_with_beta() {
+        let loss_fn = DpoLoss::with_beta(0.5);
+        assert!(approx_eq(loss_fn.config.beta, 0.5, EPS));
+        assert_eq!(loss_fn.config.variant, DpoVariant::Dpo);
+    }
+
+    #[test]
+    fn test_ipo_ctor() {
+        let loss_fn = DpoLoss::ipo(0.2);
+        assert!(approx_eq(loss_fn.config.beta, 0.2, EPS));
+        assert_eq!(loss_fn.config.variant, DpoVariant::Ipo);
+        assert!(approx_eq(loss_fn.config.label_smoothing, 0.0, EPS));
+    }
+
+    // ------------------------------------------------------------------
+    // log_sigmoid
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_log_sigmoid_zero() {
+        // log σ(0) = log(0.5) ≈ -0.693147
+        let expected = (0.5_f32).ln();
+        assert!(
+            approx_eq(DpoLoss::log_sigmoid(0.0), expected, EPS),
+            "log_sigmoid(0) = {}, expected {}",
+            DpoLoss::log_sigmoid(0.0),
+            expected
+        );
+    }
+
+    #[test]
+    fn test_log_sigmoid_large_positive() {
+        // log σ(100) ≈ 0 (σ saturates to 1)
+        let v = DpoLoss::log_sigmoid(100.0);
+        assert!(v > -1e-4 && v <= 0.0, "log_sigmoid(100) should be ≈0, got {}", v);
+    }
+
+    #[test]
+    fn test_log_sigmoid_large_negative() {
+        // log σ(x) ≈ x for large negative x (since σ(x) ≈ exp(x))
+        let v = DpoLoss::log_sigmoid(-100.0);
+        assert!(
+            approx_eq(v, -100.0, 1e-2),
+            "log_sigmoid(-100) should be ≈-100, got {}",
+            v
+        );
+    }
+
+    #[test]
+    fn test_log_sigmoid_numerically_stable() {
+        // Must not produce ±inf or NaN at extreme values
+        for &x in &[50.0_f32, -50.0_f32, 500.0, -500.0] {
+            let v = DpoLoss::log_sigmoid(x);
+            assert!(v.is_finite(), "log_sigmoid({}) produced non-finite value: {}", x, v);
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // implicit_reward
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_implicit_reward() {
+        // β=0.1, policy=-2.0, ref=-3.0  →  0.1 * (-2 - (-3)) = 0.1
+        let r = DpoLoss::implicit_reward(0.1, -2.0, -3.0);
+        assert!(approx_eq(r, 0.1, EPS), "expected 0.1, got {}", r);
+    }
+
+    // ------------------------------------------------------------------
+    // compute — DPO variant
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_compute_chosen_better() {
+        // chosen log-ratio > rejected log-ratio → positive margin → loss < log(2)
+        // policy_logp_chosen=-1, ref_logp_chosen=-3  → log_ratio_chosen = 2
+        // policy_logp_rejected=-4, ref_logp_rejected=-2 → log_ratio_rejected = -2
+        // margin = β*(2-(-2)) = 0.1*4 = 0.4
+        let loss_fn = DpoLoss::with_beta(0.1);
+        let result = loss_fn.compute(-1.0, -4.0, -3.0, -2.0);
+        let log2 = (2.0_f32).ln();
+        assert!(result.loss < log2, "loss={} should be < log(2)={}", result.loss, log2);
+        assert!(result.chosen_preferred);
+        assert!(result.reward_margin > 0.0);
+    }
+
+    #[test]
+    fn test_compute_rejected_better() {
+        // rejected has higher log-ratio → negative margin → loss > log(2)
+        // policy_logp_chosen=-4, ref_logp_chosen=-2  → log_ratio_chosen = -2
+        // policy_logp_rejected=-1, ref_logp_rejected=-3 → log_ratio_rejected = 2
+        // margin = β*(-2-2) = -0.4
+        let loss_fn = DpoLoss::with_beta(0.1);
+        let result = loss_fn.compute(-4.0, -1.0, -2.0, -3.0);
+        let log2 = (2.0_f32).ln();
+        assert!(result.loss > log2, "loss={} should be > log(2)={}", result.loss, log2);
+        assert!(!result.chosen_preferred);
+        assert!(result.reward_margin < 0.0);
+    }
+
+    #[test]
+    fn test_compute_equal() {
+        // equal log-ratios → margin=0 → loss = -log_sigmoid(0) = log(2) ≈ 0.6931
+        let loss_fn = DpoLoss::with_beta(0.1);
+        let result = loss_fn.compute(-2.0, -2.0, -2.0, -2.0);
+        let log2 = (2.0_f32).ln();
+        assert!(
+            approx_eq(result.loss, log2, 1e-4),
+            "equal log-ratios: loss={}, expected log(2)={}",
+            result.loss,
+            log2
+        );
+        assert!(approx_eq(result.reward_margin, 0.0, EPS));
+    }
+
+    #[test]
+    fn test_compute_preferred_flag() {
+        let loss_fn = DpoLoss::with_beta(0.1);
+
+        // chosen reward > rejected → true
+        let r1 = loss_fn.compute(-1.0, -5.0, -2.0, -3.0);
+        // log_ratio_chosen = -1-(-2)=1, log_ratio_rejected=-5-(-3)=-2
+        // reward_chosen = 0.1, reward_rejected = -0.2
+        assert!(r1.chosen_preferred);
+
+        // swap: now rejected is better
+        let r2 = loss_fn.compute(-5.0, -1.0, -3.0, -2.0);
+        assert!(!r2.chosen_preferred);
+    }
+
+    // ------------------------------------------------------------------
+    // IPO variant
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_ipo_loss_formula() {
+        // β=0.5, log_ratio_chosen=2.0, log_ratio_rejected=1.0
+        // target = 1/(2*0.5) = 1.0
+        // diff = 2-1-1 = 0  → loss = 0
+        // Use explicit policy/ref values that give those log-ratios:
+        //   policy_chosen=-1, ref_chosen=-3   → lrc=2
+        //   policy_rejected=-2, ref_rejected=-3 → lrr=1
+        let loss_fn = DpoLoss::ipo(0.5);
+        let result = loss_fn.compute(-1.0, -2.0, -3.0, -3.0);
+        // diff = (2 - 1) - 1/(2*0.5) = 1 - 1 = 0  → loss = 0
+        assert!(approx_eq(result.loss, 0.0, EPS), "expected 0.0, got {}", result.loss);
+    }
+
+    #[test]
+    fn test_ipo_zero_at_optimum() {
+        // At the IPO optimum: log_ratio_chosen - log_ratio_rejected == 1/(2β)
+        // β=0.1 → target = 5.0
+        // lrc=6.0, lrr=1.0 → diff=5.0-5.0=0
+        // policy_chosen=-0, ref_chosen=-6  → lrc=6
+        // policy_rejected=-1, ref_rejected=-2 → lrr=1
+        let loss_fn = DpoLoss::ipo(0.1);
+        let result = loss_fn.compute(0.0, -1.0, -6.0, -2.0);
+        assert!(approx_eq(result.loss, 0.0, EPS), "IPO at optimum: expected 0, got {}", result.loss);
+    }
+
+    // ------------------------------------------------------------------
+    // Label smoothing (Robust DPO)
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_label_smoothing_increases_loss_when_chosen_better() {
+        // When chosen is strongly preferred, label smoothing should push loss up
+        let plain = DpoLoss::new(DpoConfig { beta: 0.1, variant: DpoVariant::Dpo, label_smoothing: 0.0 });
+        let smooth = DpoLoss::new(DpoConfig { beta: 0.1, variant: DpoVariant::Dpo, label_smoothing: 0.1 });
+
+        // chosen clearly preferred: lrc=5, lrr=-5 → margin = 0.1*10 = 1.0
+        let r_plain  = plain.compute(0.0, -10.0, -5.0, -5.0);
+        let r_smooth = smooth.compute(0.0, -10.0, -5.0, -5.0);
+
+        assert!(
+            r_smooth.loss > r_plain.loss,
+            "smoothed loss ({}) should exceed plain loss ({}) when chosen is preferred",
+            r_smooth.loss,
+            r_plain.loss
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // DpoAccumulator
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_accumulator_add_single() {
+        let loss_fn = DpoLoss::with_beta(0.1);
+        let result = loss_fn.compute(-2.0, -5.0, -3.0, -3.0);
+        let mut acc = DpoAccumulator::new();
+        acc.add(&result);
+        assert!(approx_eq(acc.avg_loss(), result.loss, EPS));
+        assert_eq!(acc.n_samples(), 1);
+    }
+
+    #[test]
+    fn test_accumulator_avg_margin() {
+        let loss_fn = DpoLoss::with_beta(0.1);
+        // sample 1: lrc=1, lrr=0 → margin = 0.1
+        let r1 = loss_fn.compute(-1.0, -3.0, -2.0, -3.0);
+        // sample 2: lrc=-1, lrr=0 → margin = -0.1
+        let r2 = loss_fn.compute(-3.0, -1.0, -2.0, -1.0);
+
+        let mut acc = DpoAccumulator::new();
+        acc.add(&r1);
+        acc.add(&r2);
+
+        let expected_avg_margin = (r1.reward_margin + r2.reward_margin) / 2.0;
+        assert!(
+            approx_eq(acc.avg_reward_margin(), expected_avg_margin, EPS),
+            "avg_margin={} expected={}",
+            acc.avg_reward_margin(),
+            expected_avg_margin
+        );
+        assert_eq!(acc.n_samples(), 2);
+    }
+
+    #[test]
+    fn test_accumulator_preference_accuracy() {
+        let loss_fn = DpoLoss::with_beta(0.1);
+
+        // preferred: lrc > lrr
+        let good = loss_fn.compute(-1.0, -5.0, -2.0, -3.0);
+        assert!(good.chosen_preferred);
+
+        // not preferred: lrc < lrr
+        let bad = loss_fn.compute(-5.0, -1.0, -3.0, -2.0);
+        assert!(!bad.chosen_preferred);
+
+        // another preferred
+        let good2 = loss_fn.compute(-0.5, -4.0, -2.0, -3.0);
+        assert!(good2.chosen_preferred);
+
+        let mut acc = DpoAccumulator::new();
+        acc.add(&good);
+        acc.add(&bad);
+        acc.add(&good2);
+
+        // 2 out of 3 preferred → accuracy = 2/3
+        let expected = 2.0_f32 / 3.0;
+        assert!(
+            approx_eq(acc.preference_accuracy(), expected, EPS),
+            "accuracy={} expected={}",
+            acc.preference_accuracy(),
+            expected
+        );
+    }
+
+    #[test]
+    fn test_accumulator_reset() {
+        let loss_fn = DpoLoss::with_beta(0.1);
+        let result = loss_fn.compute(-2.0, -5.0, -3.0, -3.0);
+        let mut acc = DpoAccumulator::new();
+        acc.add(&result);
+        assert_eq!(acc.n_samples(), 1);
+
+        acc.reset();
+        assert_eq!(acc.n_samples(), 0);
+        assert!(approx_eq(acc.avg_loss(), 0.0, EPS));
+        assert!(approx_eq(acc.preference_accuracy(), 0.0, EPS));
+    }
+
+    // ------------------------------------------------------------------
+    // Batch API
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_batch_avg_loss() {
+        let loss_fn = DpoLoss::with_beta(0.1);
+
+        let pc = [-1.0_f32, -3.0, -0.5];
+        let pr = [-4.0_f32, -0.5, -3.0];
+        let rc = [-2.0_f32, -2.0, -1.0];
+        let rr = [-3.0_f32, -1.0, -2.0];
+
+        // Compute individual losses and average manually
+        let manual_avg: f32 = {
+            let mut sum = 0.0_f32;
+            for i in 0..3 {
+                sum += loss_fn.compute(pc[i], pr[i], rc[i], rr[i]).loss;
+            }
+            sum / 3.0
+        };
+
+        let (batch_avg, acc) = loss_fn.compute_batch(&pc, &pr, &rc, &rr);
+
+        assert!(
+            approx_eq(batch_avg, manual_avg, EPS),
+            "batch_avg={} manual_avg={}",
+            batch_avg,
+            manual_avg
+        );
+        assert_eq!(acc.n_samples(), 3);
+        assert!(approx_eq(acc.avg_loss(), manual_avg, EPS));
+    }
+}

crates/training/rtx-transformers/src/training/loss_functions.rs

@@ -0,0 +1,795 @@
+//! Loss functions for training: label smoothing cross-entropy, focal loss, binary variants.
+//!
+//! # References
+//! - Szegedy et al. 2016, "Rethinking the Inception Architecture" (label smoothing)
+//! - Lin et al. 2017, "Focal Loss for Dense Object Detection", arXiv:1708.02002
+
+/// How to reduce a batch of per-sample losses into a scalar.
+#[derive(Debug, Clone, PartialEq)]
+pub enum Reduction {
+    /// Arithmetic mean over the batch.
+    Mean,
+    /// Sum over the batch.
+    Sum,
+    /// No reduction — `LossResult::reduced` is set to `f32::NAN`.
+    None,
+}
+
+/// Carries both the unreduced per-sample losses and the scalar summary.
+#[derive(Debug, Clone)]
+pub struct LossResult {
+    /// One loss value per sample in the batch.
+    pub per_sample: Vec<f32>,
+    /// Reduced scalar (mean, sum, or `f32::NAN` for `Reduction::None`).
+    pub reduced: f32,
+}
+
+impl LossResult {
+    /// Arithmetic mean of per-sample losses.
+    #[inline]
+    pub fn mean(&self) -> f32 {
+        if self.per_sample.is_empty() {
+            return 0.0;
+        }
+        self.per_sample.iter().sum::<f32>() / self.per_sample.len() as f32
+    }
+
+    /// Sum of per-sample losses.
+    #[inline]
+    pub fn sum(&self) -> f32 {
+        self.per_sample.iter().sum()
+    }
+}
+
+/// Configuration for label-smoothed cross-entropy.
+#[derive(Debug, Clone)]
+pub struct LabelSmoothingConfig {
+    /// Smoothing factor ε ∈ [0, 1]. Default 0.1.
+    pub epsilon: f32,
+    /// Number of classes K (vocabulary size).
+    pub num_classes: usize,
+    /// Batch reduction strategy.
+    pub reduction: Reduction,
+}
+
+impl Default for LabelSmoothingConfig {
+    fn default() -> Self {
+        Self {
+            epsilon: 0.1,
+            num_classes: 0,
+            reduction: Reduction::Mean,
+        }
+    }
+}
+
+/// Configuration for focal loss.
+#[derive(Debug, Clone)]
+pub struct FocalLossConfig {
+    /// Focusing parameter γ ≥ 0. Default 2.0.
+    /// γ = 0 reduces to standard cross-entropy (with optional alpha weighting).
+    pub gamma: f32,
+    /// Per-class weighting factors α. `None` → α_t = 1.0 for all classes.
+    pub alpha: Option<Vec<f32>>,
+    /// Batch reduction strategy.
+    pub reduction: Reduction,
+}
+
+impl Default for FocalLossConfig {
+    fn default() -> Self {
+        Self {
+            gamma: 2.0,
+            alpha: None,
+            reduction: Reduction::Mean,
+        }
+    }
+}
+
+/// Namespace for all loss function implementations.
+pub struct LossFunctions;
+
+impl LossFunctions {
+    // -----------------------------------------------------------------------
+    // Numerically stable primitives
+    // -----------------------------------------------------------------------
+
+    /// Log-softmax computed with the max-subtraction trick for numerical stability.
+    ///
+    /// # Panics
+    /// Panics if `logits` is empty.
+    fn log_softmax(logits: &[f32]) -> Vec<f32> {
+        assert!(!logits.is_empty(), "log_softmax: logits must not be empty");
+        let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+        let shifted: Vec<f32> = logits.iter().map(|x| x - max).collect();
+        let log_sum_exp = shifted.iter().map(|x| x.exp()).sum::<f32>().ln();
+        shifted.iter().map(|x| x - log_sum_exp).collect()
+    }
+
+    /// Softmax with the max-subtraction trick.
+    ///
+    /// # Panics
+    /// Panics if `logits` is empty.
+    fn softmax(logits: &[f32]) -> Vec<f32> {
+        assert!(!logits.is_empty(), "softmax: logits must not be empty");
+        let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+        let exps: Vec<f32> = logits.iter().map(|x| (x - max).exp()).collect();
+        let sum: f32 = exps.iter().sum();
+        exps.iter().map(|e| e / sum).collect()
+    }
+
+    /// Apply reduction to a vector of per-sample losses.
+    fn apply_reduction(values: Vec<f32>, reduction: &Reduction) -> LossResult {
+        let reduced = match reduction {
+            Reduction::Mean => {
+                if values.is_empty() {
+                    0.0
+                } else {
+                    values.iter().sum::<f32>() / values.len() as f32
+                }
+            }
+            Reduction::Sum => values.iter().sum(),
+            Reduction::None => f32::NAN,
+        };
+        LossResult {
+            per_sample: values,
+            reduced,
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // 1. Standard cross-entropy
+    // -----------------------------------------------------------------------
+
+    /// Standard cross-entropy for a single sample.
+    ///
+    /// ```
+    /// CE(logits, target) = -log_softmax(logits)[target]
+    /// ```
+    ///
+    /// # Panics
+    /// Panics if `target >= logits.len()`.
+    pub fn cross_entropy(logits: &[f32], target: usize) -> f32 {
+        assert!(
+            target < logits.len(),
+            "cross_entropy: target {} out of range for {} classes",
+            target,
+            logits.len()
+        );
+        let log_probs = Self::log_softmax(logits);
+        -log_probs[target]
+    }
+
+    /// Batch cross-entropy.
+    ///
+    /// `logits` is row-major `[batch × num_classes]` (length = `batch * num_classes`).
+    /// `targets` contains the class index per sample.
+    ///
+    /// # Panics
+    /// Panics if `logits.len() != targets.len() * num_classes` or any target is out of range.
+    pub fn cross_entropy_batch(
+        logits: &[f32],
+        targets: &[usize],
+        num_classes: usize,
+        reduction: Reduction,
+    ) -> LossResult {
+        assert_eq!(
+            logits.len(),
+            targets.len() * num_classes,
+            "cross_entropy_batch: logits length mismatch"
+        );
+        let per_sample: Vec<f32> = targets
+            .iter()
+            .enumerate()
+            .map(|(i, &t)| {
+                let row = &logits[i * num_classes..(i + 1) * num_classes];
+                Self::cross_entropy(row, t)
+            })
+            .collect();
+        Self::apply_reduction(per_sample, &reduction)
+    }
+
+    // -----------------------------------------------------------------------
+    // 2. Label-smoothed cross-entropy (Szegedy et al. 2016)
+    // -----------------------------------------------------------------------
+
+    /// Label-smoothed cross-entropy for a single sample.
+    ///
+    /// ```text
+    /// loss = (1 - ε) * CE(logits, target) + ε * (−mean(log_softmax(logits)))
+    /// ```
+    ///
+    /// When ε = 0 this equals standard CE. When ε = 1 this equals the negative
+    /// mean of log-softmax (uniform target distribution).
+    ///
+    /// # Panics
+    /// Panics if `target >= logits.len()`.
+    pub fn label_smoothing_ce(logits: &[f32], target: usize, epsilon: f32) -> f32 {
+        assert!(
+            target < logits.len(),
+            "label_smoothing_ce: target {} out of range for {} classes",
+            target,
+            logits.len()
+        );
+        let log_probs = Self::log_softmax(logits);
+        let k = logits.len() as f32;
+        let hard_loss = -log_probs[target];
+        // Negative mean of log-softmax = cross-entropy against uniform distribution.
+        let smooth_term = -log_probs.iter().sum::<f32>() / k;
+        (1.0 - epsilon) * hard_loss + epsilon * smooth_term
+    }
+
+    /// Batch label-smoothed cross-entropy.
+    ///
+    /// `logits` is row-major `[batch × num_classes]`.
+    pub fn label_smoothing_ce_batch(
+        logits: &[f32],
+        targets: &[usize],
+        num_classes: usize,
+        config: &LabelSmoothingConfig,
+    ) -> LossResult {
+        assert_eq!(
+            logits.len(),
+            targets.len() * num_classes,
+            "label_smoothing_ce_batch: logits length mismatch"
+        );
+        let per_sample: Vec<f32> = targets
+            .iter()
+            .enumerate()
+            .map(|(i, &t)| {
+                let row = &logits[i * num_classes..(i + 1) * num_classes];
+                Self::label_smoothing_ce(row, t, config.epsilon)
+            })
+            .collect();
+        Self::apply_reduction(per_sample, &config.reduction)
+    }
+
+    // -----------------------------------------------------------------------
+    // 3. Focal loss (Lin et al. 2017, arXiv:1708.02002)
+    // -----------------------------------------------------------------------
+
+    /// Focal loss for a single sample.
+    ///
+    /// ```text
+    /// p_t      = softmax(logits)[target]
+    /// alpha_t  = alpha[target]  (or 1.0 if alpha is None)
+    /// loss     = -alpha_t * (1 - p_t)^gamma * log(p_t)
+    /// ```
+    ///
+    /// γ = 0 with α = None reduces to standard cross-entropy.
+    ///
+    /// # Panics
+    /// Panics if `target >= logits.len()` or if `alpha` has fewer entries than `logits`.
+    pub fn focal_loss(logits: &[f32], target: usize, config: &FocalLossConfig) -> f32 {
+        assert!(
+            target < logits.len(),
+            "focal_loss: target {} out of range for {} classes",
+            target,
+            logits.len()
+        );
+        if let Some(ref a) = config.alpha {
+            assert!(
+                a.len() >= logits.len(),
+                "focal_loss: alpha length {} < num_classes {}",
+                a.len(),
+                logits.len()
+            );
+        }
+        let probs = Self::softmax(logits);
+        let log_probs = Self::log_softmax(logits);
+        let p_t = probs[target];
+        let alpha_t = config.alpha.as_ref().map_or(1.0, |a| a[target]);
+        let focal_weight = (1.0 - p_t).powf(config.gamma);
+        -alpha_t * focal_weight * log_probs[target]
+    }
+
+    /// Batch focal loss.
+    ///
+    /// `logits` is row-major `[batch × num_classes]`.
+    pub fn focal_loss_batch(
+        logits: &[f32],
+        targets: &[usize],
+        num_classes: usize,
+        config: &FocalLossConfig,
+    ) -> LossResult {
+        assert_eq!(
+            logits.len(),
+            targets.len() * num_classes,
+            "focal_loss_batch: logits length mismatch"
+        );
+        let per_sample: Vec<f32> = targets
+            .iter()
+            .enumerate()
+            .map(|(i, &t)| {
+                let row = &logits[i * num_classes..(i + 1) * num_classes];
+                Self::focal_loss(row, t, config)
+            })
+            .collect();
+        Self::apply_reduction(per_sample, &config.reduction)
+    }
+
+    // -----------------------------------------------------------------------
+    // 4. Smoothed focal loss (label smoothing + focal modulation)
+    // -----------------------------------------------------------------------
+
+    /// Smoothed focal loss: apply label smoothing first, then focal modulation.
+    ///
+    /// ```text
+    /// p_t          = softmax(logits)[target]
+    /// log_p_smooth = log((1-ε)*one_hot[target] + ε/K) ≈ log_softmax shifted by smoothing
+    /// loss         = -alpha_t * (1 - p_t)^gamma * log_p_smooth[target]
+    /// ```
+    ///
+    /// The focal weight uses the raw softmax probability p_t so that
+    /// well-classified examples are still down-weighted even with smoothing.
+    ///
+    /// # Panics
+    /// Panics if `target >= logits.len()`.
+    pub fn smoothed_focal_loss(
+        logits: &[f32],
+        target: usize,
+        smoothing_config: &LabelSmoothingConfig,
+        focal_config: &FocalLossConfig,
+    ) -> f32 {
+        assert!(
+            target < logits.len(),
+            "smoothed_focal_loss: target {} out of range for {} classes",
+            target,
+            logits.len()
+        );
+        let probs = Self::softmax(logits);
+        let log_probs = Self::log_softmax(logits);
+        let k = logits.len() as f32;
+        let epsilon = smoothing_config.epsilon;
+
+        // Compute smooth log-prob at the target index.
+        // p_smooth[target] = (1-ε)*1 + ε/K  (one-hot at target)
+        // p_smooth[other]  = ε/K
+        // We compute it directly from log_probs to avoid a second softmax.
+        // p_smooth[target] = (1-ε)*exp(log_probs[target]) + ε/K  — but that requires
+        // re-computing log of a mixture, so we do it explicitly.
+        let p_raw_target = probs[target];
+        let p_smooth_target = (1.0 - epsilon) * p_raw_target + epsilon / k;
+        let log_p_smooth_target = p_smooth_target.ln();
+        // Ignore the tiny discrepancy between log_probs[target] and log(p_smooth[target]);
+        // use the true smoothed value.
+        let _ = log_probs; // log_probs computed above; referenced only for documentation clarity.
+
+        let p_t = p_raw_target;
+        let alpha_t = focal_config.alpha.as_ref().map_or(1.0, |a| a[target]);
+        let focal_weight = (1.0 - p_t).powf(focal_config.gamma);
+        -alpha_t * focal_weight * log_p_smooth_target
+    }
+
+    // -----------------------------------------------------------------------
+    // 5. Binary cross-entropy (numerically stable)
+    // -----------------------------------------------------------------------
+
+    /// Numerically stable binary cross-entropy from a raw logit.
+    ///
+    /// ```text
+    /// BCE(logit, label) = max(logit, 0) - logit * label + log(1 + exp(-|logit|))
+    /// ```
+    ///
+    /// This form avoids overflow for large positive logits and underflow for
+    /// large negative logits.
+    ///
+    /// `label` should be 0.0 or 1.0.
+    pub fn binary_ce(logit: f32, label: f32) -> f32 {
+        logit.max(0.0) - logit * label + (1.0 + (-logit.abs()).exp()).ln()
+    }
+
+    // -----------------------------------------------------------------------
+    // 6. Binary focal loss
+    // -----------------------------------------------------------------------
+
+    /// Binary focal loss for a single logit/label pair.
+    ///
+    /// ```text
+    /// p    = sigmoid(logit)
+    /// p_t  = p  if label == 1 else 1 - p
+    /// loss = -alpha_t * (1 - p_t)^gamma * log(p_t)
+    /// ```
+    ///
+    /// `alpha` is the weighting factor for the positive class (label = 1).
+    /// For the negative class, the complement `(1 - alpha)` is used.
+    /// `gamma = 0, alpha = 1.0` reproduces standard binary cross-entropy.
+    pub fn binary_focal(logit: f32, label: f32, gamma: f32, alpha: f32) -> f32 {
+        let p = sigmoid(logit);
+        let (p_t, alpha_t) = if label >= 0.5 {
+            (p, alpha)
+        } else {
+            (1.0 - p, 1.0 - alpha)
+        };
+        // Clamp to avoid log(0).
+        let p_t_clamped = p_t.clamp(1e-7, 1.0 - 1e-7);
+        let focal_weight = (1.0 - p_t).powf(gamma);
+        -alpha_t * focal_weight * p_t_clamped.ln()
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Private utility
+// ---------------------------------------------------------------------------
+
+/// Sigmoid function.
+#[inline]
+fn sigmoid(x: f32) -> f32 {
+    1.0 / (1.0 + (-x).exp())
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Tolerance for floating-point comparisons.
+    const EPS: f32 = 1e-5;
+
+    fn approx_eq(a: f32, b: f32, tol: f32) -> bool {
+        (a - b).abs() < tol
+    }
+
+    // -----------------------------------------------------------------------
+    // Standard cross-entropy
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_cross_entropy_uniform() {
+        // For uniform logits [0, 0, 0, 0] each class has prob 1/4.
+        // CE = log(4) ≈ 1.386294
+        let logits = vec![0.0f32; 4];
+        let ce = LossFunctions::cross_entropy(&logits, 0);
+        let expected = (4.0f32).ln();
+        assert!(
+            approx_eq(ce, expected, EPS),
+            "uniform CE: got {ce}, expected {expected}"
+        );
+    }
+
+    #[test]
+    fn test_cross_entropy_confident() {
+        // Very high logit on the correct class → CE ≈ 0.
+        let mut logits = vec![-100.0f32; 4];
+        logits[2] = 100.0;
+        let ce = LossFunctions::cross_entropy(&logits, 2);
+        assert!(ce < 1e-3, "confident CE should be near 0, got {ce}");
+    }
+
+    #[test]
+    fn test_cross_entropy_batch_mean() {
+        // Two samples, two classes. Manual reference values.
+        let logits = vec![
+            1.0f32, 0.0, // sample 0: target 0
+            0.0, 1.0,    // sample 1: target 1
+        ];
+        let targets = vec![0usize, 1];
+        let result =
+            LossFunctions::cross_entropy_batch(&logits, &targets, 2, Reduction::Mean);
+        // Each individual CE should be the same (symmetric).
+        let ce0 = LossFunctions::cross_entropy(&logits[0..2], 0);
+        let ce1 = LossFunctions::cross_entropy(&logits[2..4], 1);
+        let expected_mean = (ce0 + ce1) / 2.0;
+        assert!(
+            approx_eq(result.reduced, expected_mean, EPS),
+            "batch mean: got {}, expected {}",
+            result.reduced,
+            expected_mean
+        );
+    }
+
+    #[test]
+    fn test_cross_entropy_batch_sum() {
+        let logits = vec![1.0f32, 0.0, 0.0, 1.0];
+        let targets = vec![0usize, 1];
+        let result =
+            LossFunctions::cross_entropy_batch(&logits, &targets, 2, Reduction::Sum);
+        let ce0 = LossFunctions::cross_entropy(&logits[0..2], 0);
+        let ce1 = LossFunctions::cross_entropy(&logits[2..4], 1);
+        let expected_sum = ce0 + ce1;
+        assert!(
+            approx_eq(result.reduced, expected_sum, EPS),
+            "batch sum: got {}, expected {}",
+            result.reduced,
+            expected_sum
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Label smoothing
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_label_smoothing_less_than_hard_ce_when_confident() {
+        // With a very confident correct prediction, smoothed CE > hard CE
+        // because smoothing pulls probability mass toward the uniform distribution,
+        // introducing penalty for over-confidence.
+        let mut logits = vec![-10.0f32; 4];
+        logits[0] = 10.0;
+        let hard_ce = LossFunctions::cross_entropy(&logits, 0);
+        let smooth_ce = LossFunctions::label_smoothing_ce(&logits, 0, 0.1);
+        assert!(
+            smooth_ce > hard_ce,
+            "smoothed CE ({smooth_ce}) should exceed hard CE ({hard_ce}) when model is overconfident"
+        );
+    }
+
+    #[test]
+    fn test_label_smoothing_epsilon_zero() {
+        let logits = vec![2.0f32, 1.0, 0.5, -0.5];
+        let hard_ce = LossFunctions::cross_entropy(&logits, 1);
+        let smooth_ce = LossFunctions::label_smoothing_ce(&logits, 1, 0.0);
+        assert!(
+            approx_eq(hard_ce, smooth_ce, EPS),
+            "ε=0 smoothed CE ({smooth_ce}) must equal hard CE ({hard_ce})"
+        );
+    }
+
+    #[test]
+    fn test_label_smoothing_epsilon_one() {
+        // ε=1: smooth_term = -mean(log_softmax(logits))
+        let logits = vec![2.0f32, 1.0, 0.5, -0.5];
+        let log_probs = {
+            let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+            let shifted: Vec<f32> = logits.iter().map(|x| x - max).collect();
+            let lse = shifted.iter().map(|x| x.exp()).sum::<f32>().ln();
+            shifted.iter().map(|x| x - lse).collect::<Vec<_>>()
+        };
+        let expected = -log_probs.iter().sum::<f32>() / log_probs.len() as f32;
+        let smooth_ce = LossFunctions::label_smoothing_ce(&logits, 0, 1.0);
+        assert!(
+            approx_eq(smooth_ce, expected, EPS),
+            "ε=1 CE ({smooth_ce}) must equal -mean(log_softmax) ({expected})"
+        );
+    }
+
+    #[test]
+    fn test_label_smoothing_batch() {
+        let logits = vec![1.0f32, 0.0, 0.0, 1.0, 0.5, 0.5];
+        let targets = vec![0usize, 1, 0];
+        let config = LabelSmoothingConfig {
+            epsilon: 0.1,
+            num_classes: 2,
+            reduction: Reduction::Mean,
+        };
+        let result =
+            LossFunctions::label_smoothing_ce_batch(&logits, &targets, 2, &config);
+        assert_eq!(result.per_sample.len(), 3, "batch must yield 3 per-sample losses");
+        // Verify reduced == mean of per-sample.
+        let expected_mean = result.per_sample.iter().sum::<f32>() / 3.0;
+        assert!(
+            approx_eq(result.reduced, expected_mean, EPS),
+            "batch reduced ({}) != mean of per_sample ({})",
+            result.reduced,
+            expected_mean
+        );
+    }
+
+    #[test]
+    fn test_label_smoothing_reduces_overconfidence() {
+        // Measure the entropy of the model's distribution under hard vs smoothed targets.
+        // After gradient descent, smoothed training keeps entropy higher.
+        // As a proxy: with ε > 0 the smoothed loss is strictly larger when the model is
+        // very confident, which forces the optimizer to regularize confidence.
+        let mut logits = vec![-20.0f32; 10];
+        logits[3] = 20.0;
+        let hard = LossFunctions::cross_entropy(&logits, 3);
+        let smooth = LossFunctions::label_smoothing_ce(&logits, 3, 0.2);
+        // Smoothed loss should be larger (penalizing overconfidence).
+        assert!(
+            smooth > hard,
+            "smoothed loss ({smooth}) should exceed hard loss ({hard}) for overconfident model"
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Focal loss
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_focal_gamma_zero_equals_ce() {
+        let logits = vec![2.0f32, 1.0, 0.5, -0.5];
+        let ce = LossFunctions::cross_entropy(&logits, 2);
+        let config = FocalLossConfig {
+            gamma: 0.0,
+            alpha: None,
+            reduction: Reduction::Mean,
+        };
+        let fl = LossFunctions::focal_loss(&logits, 2, &config);
+        assert!(
+            approx_eq(fl, ce, EPS),
+            "focal(γ=0, α=None) ({fl}) must equal CE ({ce})"
+        );
+    }
+
+    #[test]
+    fn test_focal_high_gamma_downweights_easy() {
+        // An easy example has high p_t → (1 - p_t)^γ is small.
+        // Compare two samples: one easy (correct class has high logit) and one hard.
+        let logits_easy = vec![10.0f32, -10.0]; // p_0 ≈ 1.0 → very easy
+        let logits_hard = vec![0.1f32, 0.0];    // p_0 ≈ 0.52 → hard
+
+        let config = FocalLossConfig {
+            gamma: 2.0,
+            alpha: None,
+            reduction: Reduction::Mean,
+        };
+        let fl_easy = LossFunctions::focal_loss(&logits_easy, 0, &config);
+        let fl_hard = LossFunctions::focal_loss(&logits_hard, 0, &config);
+        assert!(
+            fl_easy < fl_hard,
+            "easy focal loss ({fl_easy}) should be smaller than hard focal loss ({fl_hard})"
+        );
+    }
+
+    #[test]
+    fn test_focal_alpha_weighting() {
+        let logits = vec![1.0f32, 0.0, 0.0];
+        let config_no_alpha = FocalLossConfig {
+            gamma: 0.0,
+            alpha: None,
+            reduction: Reduction::Mean,
+        };
+        let config_alpha = FocalLossConfig {
+            gamma: 0.0,
+            alpha: Some(vec![0.25, 1.0, 1.0]),
+            reduction: Reduction::Mean,
+        };
+        let fl_no_alpha = LossFunctions::focal_loss(&logits, 0, &config_no_alpha);
+        let fl_alpha = LossFunctions::focal_loss(&logits, 0, &config_alpha);
+        // With α_0 = 0.25 and γ = 0: fl_alpha should be 0.25 * fl_no_alpha.
+        assert!(
+            approx_eq(fl_alpha, 0.25 * fl_no_alpha, EPS),
+            "alpha=0.25 should scale loss: got {fl_alpha}, expected {}",
+            0.25 * fl_no_alpha
+        );
+    }
+
+    #[test]
+    fn test_focal_batch_mean() {
+        let logits = vec![2.0f32, 1.0, 0.0, 0.0, 2.0, 1.0];
+        let targets = vec![0usize, 0, 1];
+        let config = FocalLossConfig {
+            gamma: 2.0,
+            alpha: None,
+            reduction: Reduction::Mean,
+        };
+        let result = LossFunctions::focal_loss_batch(&logits, &targets, 2, &config);
+        let expected_mean = result.per_sample.iter().sum::<f32>() / 3.0;
+        assert!(
+            approx_eq(result.reduced, expected_mean, EPS),
+            "batch mean focal: got {}, expected {}",
+            result.reduced,
+            expected_mean
+        );
+    }
+
+    #[test]
+    fn test_focal_batch_reduction_none() {
+        let logits = vec![1.0f32, 0.0, 0.0, 1.0];
+        let targets = vec![0usize, 1];
+        let config = FocalLossConfig {
+            gamma: 1.0,
+            alpha: None,
+            reduction: Reduction::None,
+        };
+        let result = LossFunctions::focal_loss_batch(&logits, &targets, 2, &config);
+        assert_eq!(result.per_sample.len(), 2);
+        // `reduced` must be NaN for Reduction::None.
+        assert!(
+            result.reduced.is_nan(),
+            "Reduction::None must set reduced = NaN"
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Binary cross-entropy
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_binary_ce_label_one() {
+        // BCE(0.0, 1.0): sigmoid(0) = 0.5, so -log(0.5) = log(2) ≈ 0.693147.
+        let bce = LossFunctions::binary_ce(0.0, 1.0);
+        let expected = (2.0f32).ln();
+        assert!(
+            approx_eq(bce, expected, EPS),
+            "BCE(0,1)={bce}, expected {expected}"
+        );
+    }
+
+    #[test]
+    fn test_binary_ce_label_zero() {
+        // BCE(0.0, 0.0): -log(1 - sigmoid(0)) = -log(0.5) = log(2).
+        let bce = LossFunctions::binary_ce(0.0, 0.0);
+        let expected = (2.0f32).ln();
+        assert!(
+            approx_eq(bce, expected, EPS),
+            "BCE(0,0)={bce}, expected {expected}"
+        );
+    }
+
+    #[test]
+    fn test_binary_ce_confident_correct() {
+        // BCE(10.0, 1.0): sigmoid(10) ≈ 1 → loss ≈ 0.
+        let bce = LossFunctions::binary_ce(10.0, 1.0);
+        assert!(bce < 1e-3, "BCE(10,1) should be ≈0, got {bce}");
+    }
+
+    #[test]
+    fn test_binary_ce_numerically_stable() {
+        // BCE(100.0, 0.0) must be finite — the naive -log(1-sigmoid(100)) would overflow.
+        let bce = LossFunctions::binary_ce(100.0, 0.0);
+        assert!(bce.is_finite(), "BCE(100,0) must be finite, got {bce}");
+        // Value should be approximately 100 (the logit itself dominates).
+        assert!(
+            approx_eq(bce, 100.0, 1e-3),
+            "BCE(100,0) ≈ 100, got {bce}"
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Binary focal loss
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_binary_focal_gamma_zero_equals_bce() {
+        // binary_focal(logit, label, gamma=0, alpha=1) should equal binary_ce.
+        let logit = 0.5f32;
+        let label = 1.0f32;
+        let bce = LossFunctions::binary_ce(logit, label);
+        let bfl = LossFunctions::binary_focal(logit, label, 0.0, 1.0);
+        // Allow slight tolerance due to clamping in binary_focal.
+        assert!(
+            approx_eq(bce, bfl, 1e-4),
+            "binary_focal(γ=0,α=1) ({bfl}) should equal binary_ce ({bce})"
+        );
+    }
+
+    #[test]
+    fn test_binary_focal_easy_example() {
+        // p_t = sigmoid(10) ≈ 0.9999546 → (1-p_t)^2 ≈ (4.54e-5)^2 ≈ 2e-9.
+        // focal weight should be very small.
+        let logit = 10.0f32;
+        let label = 1.0f32;
+        let p = sigmoid(logit);
+        let focal_weight = (1.0 - p).powi(2);
+        let bfl = LossFunctions::binary_focal(logit, label, 2.0, 1.0);
+        // The loss should be very small.
+        assert!(
+            bfl < 1e-4,
+            "easy example binary focal ({bfl}) should be tiny (focal_weight={focal_weight})"
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // LossResult utilities
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_loss_result_mean_and_sum() {
+        let result = LossResult {
+            per_sample: vec![1.0, 2.0, 3.0],
+            reduced: 2.0, // mean
+        };
+        assert!(approx_eq(result.mean(), 2.0, EPS), "mean mismatch");
+        assert!(approx_eq(result.sum(), 6.0, EPS), "sum mismatch");
+    }
+
+    #[test]
+    fn test_reduction_none_preserves_all() {
+        let logits = vec![1.0f32, 0.0, 0.0, 1.0, 0.5, -0.5];
+        let targets = vec![0usize, 1, 0];
+        let result =
+            LossFunctions::cross_entropy_batch(&logits, &targets, 2, Reduction::None);
+        assert_eq!(
+            result.per_sample.len(),
+            3,
+            "Reduction::None must preserve all per-sample values"
+        );
+        assert!(
+            result.reduced.is_nan(),
+            "Reduction::None must set reduced = NaN"
+        );
+    }
+}

crates/training/rtx-transformers/src/training/mod.rs

@@ -58,6 +58,14 @@ pub use model_ema::{ModelEma, ModelEmaConfig};
 pub mod swa;
 pub use swa::{SwaBuffer, SwaSchedule, SwaTrainer, SwagBuffer};
 
+pub mod dpo;
+pub use dpo::{DpoAccumulator, DpoConfig, DpoLoss, DpoLossResult, DpoVariant};
+
+pub mod loss_functions;
+pub use loss_functions::{
+    FocalLossConfig, LabelSmoothingConfig, LossFunctions, LossResult, Reduction,
+};
+
 /// Training state structure
 #[derive(Debug, Clone)]
 pub struct TrainingState {