diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl
index 53ece98..d224bc0 100644
--- a/benchmarks/titanic-logloss.jl
+++ b/benchmarks/titanic-logloss.jl
@@ -85,6 +85,12 @@ learner = NeuroTabRegressor(
     device=:cpu,
     backend=:reactant
 )
+learner.embedding_config
+# main / lux-new
+# NeuroTabModels.Models.Embeddings.EmbeddingConfig(:batchnorm, 1, :relu, 32, 32, 0.01f0)
+# mnca/temporal
+# NeuroTabModels.Models.Embeddings.EmbeddingLayer{NeuroTabModels.Models.Embeddings.BatchNormEmbeddings, Nothing}(NeuroTabModels.Models.Embeddings.BatchNormEmbeddings(), nothing)
+# NeuroTabModels.Models.Embeddings.EmbeddingLayer{NeuroTabModels.Models.Embeddings.PiecewiseLinearEmbeddings, Nothing}(NeuroTabModels.Models.Embeddings.PiecewiseLinearEmbeddings(8, 16, :identity, :B), nothing)
 
 @time m = NeuroTabModels.fit(
     learner,
diff --git a/docs/make.jl b/docs/make.jl
index fd2806d..6ad07d9 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -7,6 +7,7 @@ pages = [
     "Quick start" => "quick-start.md",
     "API" => "API.md",
     "Design" => "design.md",
+    "Embeddings" => "embeddings.md",
     "Models" => [
         "Interface" => "models/models.md",
         "MLP" => "models/mlp.md",
diff --git a/docs/src/embeddings.md b/docs/src/embeddings.md
index 5d4d933..bfdcc3b 100644
--- a/docs/src/embeddings.md
+++ b/docs/src/embeddings.md
@@ -1,5 +1,43 @@
 # Embeddings
 
+Embeddings transform raw input columns before they reach a model backbone. Numerical
+embeddings operate column-wise, while temporal embeddings reserve one column for
+time features and concatenate that branch with the remaining numerical features.
+
+## Available Embeddings
+
+- `IdentityEmbedding`: leaves numerical features unchanged.
+- `LinearEmbeddings`: projects each numerical feature to a learned vector.
+- `PeriodicEmbeddings`: expands each feature with learned sinusoidal terms before projection.
+- `PiecewiseLinearEmbeddings`: computes feature bins from the training data and embeds the resulting piecewise-linear encoding.
+- `BatchNormEmbeddings`: batch-normalizes raw numerical features without expanding them.
+- `TemporalEmbeddings`: embeds one time column with Fourier features and an optional trend term.
+
+## Configuration
+
+Use `EmbeddingLayer` to combine a numerical embedding with an optional temporal
+branch. A dictionary form is also available for model configuration dictionaries.
+
+```julia
+EmbeddingLayer(PeriodicEmbeddings(; d_embedding=24))
+
+EmbeddingLayer(;
+    num=PiecewiseLinearEmbeddings(; bins=32),
+    temp=TemporalEmbeddings(; index=1),
+)
+
+EmbeddingLayer(Dict(
+    :embedding_type => :periodic,
+    :d_embedding => 24,
+    :temporal => Dict(:index => 1),
+))
+```
+
+Piecewise-linear and temporal embeddings need training data when the embedding chain is
+built. Use `needs_x_train` to check that requirement for a spec.
+
+## API
+
 ```@autodocs
 Modules = [NeuroTabModels.Models.Embeddings]
 ```
diff --git a/ext/NeuroTabModelsReactantExt.jl b/ext/NeuroTabModelsReactantExt.jl
index 791cc74..c4e5c90 100644
--- a/ext/NeuroTabModelsReactantExt.jl
+++ b/ext/NeuroTabModelsReactantExt.jl
@@ -16,12 +16,17 @@ function _get_device(::Val{:reactant}, ::Val{D}; gpuID::Integer=0) where {D}
     return reactant_device()
 end
 
+_same_shape(a::AbstractArray, b::AbstractArray) = size(a) == size(b)
+_same_shape(a::Tuple, b::Tuple) =
+    length(a) == length(b) && all(_same_shape(ai, bi) for (ai, bi) in zip(a, b))
+_same_shape(_, _) = false
+
 function _infer_loop(::Val{:reactant}, chain, data, x0, dev, cdev, ps, st)
     compiled = @compile _forward_reduce(chain, dev(x0), ps, st)
 
     preds = Vector{AbstractArray}()
     for x in data
-        if size(x) == size(x0)
+        if _same_shape(x, x0)
             pred = compiled(chain, dev(x), ps, st)
         else
             pred = Reactant.@jit _forward_reduce(chain, dev(x), ps, st)
diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl
index 04b7f7b..5fbaef6 100644
--- a/src/Fit/callback.jl
+++ b/src/Fit/callback.jl
@@ -7,6 +7,7 @@ using ..Learners: LearnerTypes
 using ...Infer: reduce_pred, _get_device
 using ..Data: get_df_loader_train
 using ..Metrics
+using ...Models
 
 using Lux: Training, testmode
 
@@ -26,7 +27,8 @@ end
 function CallBack(
     config::LearnerTypes,
     df::AbstractDataFrame,
-    cache;
+    cache,
+    m;
     feature_names,
     target_name,
     weight_name=nothing,
@@ -43,6 +45,7 @@ function CallBack(
     deval = get_df_loader_train(dfg; feature_names, target_name, weight_name, offset_name, scalers, batchsize, shuffle=false) |> dev
 
     ps, st = ts.parameters, testmode(ts.states)
+    deval = Models.eval_dataloader(m.chain, m.info, deval, dev, ps, st)
     d0 = first(deval)
     eval_compiled = _build_eval_step(ts.model, feval, d0, ps, st; reactant=config.backend == :reactant)
 
diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl
index cf9e41b..a3beec5 100644
--- a/src/Fit/fit.jl
+++ b/src/Fit/fit.jl
@@ -9,7 +9,7 @@ using ..Losses
 using ..Metrics
 using ..Infer: reduce_pred, _get_device
 
-import Random: Xoshiro
+import Random: Xoshiro, default_rng
 import Statistics: mean, std
 import MLJModelInterface: fit
 import Optimisers: OptimiserChain, WeightDecay, NAdam, Adam
@@ -73,19 +73,12 @@ function init(
 
     # Build chain: optional embeddings + architecture backbone
     embed_config = config.embedding_config
-    if isnothing(embed_config)
-        chain = config.arch(; nfeats, outsize)
-    else
-        if embed_config.embedding_type == :piecewise
-            x_train = Matrix{Float32}(df[:, feature_names])
-        else
-            x_train = nothing
-        end
-        embed_chain = embed_config(; nfeats, x_train)
-        d_in = nfeats * embed_config.d_embedding
-        d_features = fill(embed_config.d_embedding, nfeats)
-        chain = Chain(embed_chain, config.arch(; nfeats=d_in, outsize, d_features, scaling_init_override=:normal))
-    end
+    x_train = Models.Embeddings.needs_x_train(embed_config) ? Matrix{Float32}(df[:, feature_names]) : nothing
+    embed_chain = Models.Embeddings.build_embedding_chain(embed_config, nfeats; x_train)
+    d_in = Models.Embeddings.embedding_width(embed_chain, randn(Float32, nfeats, 2), default_rng())
+    chain = Models.build_chain(config.arch, embed_chain;
+        nfeats, outsize, d_in, embedding_config=embed_config, loss_type=L)
+
 
     info = Dict(
         :nrounds => 0,
@@ -105,6 +98,8 @@ function init(
 
     rng = Xoshiro(config.seed)
     ps, st = Lux.setup(rng, m.chain) |> dev
+    data = Models.train_dataloader(config.arch, m, data, df;
+        feature_names, target_name, loss_type=L, scalers, batchsize, dev, rng)
     opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd))
     ts = Training.TrainState(m.chain, ps, st, opt)
 
@@ -164,7 +159,7 @@ function fit(
 
     logger = nothing
     if !isnothing(deval)
-        cb = CallBack(config, deval, cache; feature_names, target_name, weight_name, offset_name, group_key)
+        cb = CallBack(config, deval, cache, m; feature_names, target_name, weight_name, offset_name, group_key)
         logger = init_logger(config)
         cb(logger, 0, cache[:train_state])
         (verbosity > 0) && @info "Init training" metric = logger[:metrics][end]
diff --git a/src/learners.jl b/src/learners.jl
index 811e00a..849ea18 100644
--- a/src/learners.jl
+++ b/src/learners.jl
@@ -5,14 +5,19 @@ import MLJModelInterface: fit, update, predict, schema
 import Random
 
 using ..Models
-using ..Models: EmbeddingConfig
+using ..Models: AbstractEmbedding, IdentityEmbedding, EmbeddingLayer
 export NeuroTabRegressor, NeuroTabClassifier, LearnerTypes
 
+_to_embedding(::Nothing)            = IdentityEmbedding()
+_to_embedding(e::AbstractEmbedding) = e
+_to_embedding(d::AbstractDict)      = EmbeddingLayer(d)
+_to_embedding(x) = error("`embedding_config` must be `nothing`, an `AbstractDict`, or an `AbstractEmbedding`; got $(typeof(x)).")
+
 mutable struct NeuroTabRegressor <: MMI.Deterministic
   loss::Symbol
   metric::Symbol
   arch::Architecture
-  embedding_config::Union{Nothing,EmbeddingConfig}
+  embedding_config::AbstractEmbedding
   nrounds::Int
   early_stopping_rounds::Int
   lr::Float32
@@ -47,8 +52,10 @@ A model type for constructing a NeuroTabRegressor, based on [NeuroTabModels.jl](
 - `backend=:zygote`:        Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`.
 - `device=:gpu`:            Execution device. One of `:cpu` or `:gpu`.
 - `gpuID=0`:                GPU device to use, only relevant if `device = :gpu`. `0` auto-selects.
-- `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings.
-  E.g. `embedding_config=Dict(:embedding_type => :periodic, :d_embedding => 24)`.
+- `embedding_config=nothing`: Optional numerical/temporal embeddings. Accepts `nothing` (no-op),
+  an `AbstractEmbedding` (e.g. `EmbeddingLayer(num=PeriodicEmbeddings(d_embedding=24))`), or an
+  `AbstractDict` selecting the type via `:embedding_type` (e.g.
+  `Dict(:embedding_type => :periodic, :d_embedding => 24)`).
 
 `backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD.
 `backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models.
@@ -194,11 +201,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...)
     @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU."
   end
 
-  # Build EmbeddingConfig from Dict if needed
-  embed = args[:embedding_config]
-  if embed isa AbstractDict
-    embed = EmbeddingConfig(; embed...)
-  end
+  embed = _to_embedding(args[:embedding_config])
 
   config = NeuroTabRegressor(
     loss,
@@ -231,7 +234,7 @@ mutable struct NeuroTabClassifier <: MMI.Probabilistic
   loss::Symbol
   metric::Symbol
   arch::Architecture
-  embedding_config::Union{Nothing,EmbeddingConfig}
+  embedding_config::AbstractEmbedding
   nrounds::Int
   early_stopping_rounds::Int
   lr::Float32
@@ -259,7 +262,8 @@ A model type for constructing a NeuroTabClassifier, based on [NeuroTabModels.jl]
 - `backend=:zygote`:        Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`.
 - `device=:gpu`:            Execution device. One of `:cpu` or `:gpu`.
 - `gpuID=0`:                GPU device to use, only relevant if `device = :gpu`. `0` auto-selects.
-- `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings.
+- `embedding_config=nothing`: Optional numerical/temporal embeddings. Accepts `nothing` (no-op),
+  an `AbstractEmbedding`, or an `AbstractDict` selecting the type via `:embedding_type`.
 
 `backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD.
 `backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models.
@@ -390,11 +394,7 @@ function NeuroTabClassifier(arch::Architecture; kwargs...)
     @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU."
   end
 
-  # Build EmbeddingConfig from Dict if needed
-  embed = args[:embedding_config]
-  if embed isa AbstractDict
-    embed = EmbeddingConfig(; embed...)
-  end
+  embed = _to_embedding(args[:embedding_config])
 
   config = NeuroTabClassifier(
     :mlogloss,
diff --git a/src/models/ModernNCA/modernnca.jl b/src/models/ModernNCA/modernnca.jl
new file mode 100644
index 0000000..f60524c
--- /dev/null
+++ b/src/models/ModernNCA/modernnca.jl
@@ -0,0 +1,365 @@
+module ModernNCA
+
+export ModernNCAConfig
+
+import ..Models
+import ..Models: Architecture, NeuroTabModel
+import ..Losses: LossType, MSE, MAE, LogLoss, MLogLoss, GaussianMLE
+
+using Lux
+using LuxCore
+using NNlib: relu, softmax
+using Random: randperm, AbstractRNG
+using DataFrames: AbstractDataFrame, select
+using CategoricalArrays
+
+"""
+    ModernNCAConfig(; d_embedding=128, n_blocks=2, d_block=256,
+                     dropout=0.1, temperature=1.0, sample_rate=0.8,
+                     max_candidates=8192, eps=1f-8)
+
+Hyperparameters for ModernNCA. Pass as the `arch` argument to `NeuroTabRegressor`
+or `NeuroTabClassifier`.
+
+# Arguments
+- `d_embedding`: encoder output dimension.
+- `n_blocks`, `d_block`: number and hidden width of post-encoder MLP blocks.
+- `dropout`: dropout rate inside each block; skipped when `<= 0`.
+- `temperature`: softmax temperature on negative pairwise distances.
+- `sample_rate`: fraction of the batch complement sampled as candidates per step
+  (Stochastic Neighborhood Sampling). `1.0` uses the full complement.
+- `max_candidates`: cap on sampled training candidates per batch. Set `<= 0`
+  to disable the cap.
+- `eps`: numerical floor in `sqrt` and as a `temperature` lower bound.
+"""
+struct ModernNCAConfig <: Architecture
+    d_embedding::Int
+    n_blocks::Int
+    d_block::Int
+    dropout::Float32
+    temperature::Float32
+    sample_rate::Float32
+    max_candidates::Int
+    eps::Float32
+end
+
+"""
+    ModernNCAModel
+
+Lux wrapper holding the backbone encoder, config, output size, and loss type.
+"""
+struct ModernNCAModel{B,LT} <: LuxCore.AbstractLuxWrapperLayer{:backbone}
+    backbone::B
+    cfg::ModernNCAConfig
+    outsize::Int
+    loss_type::Type{LT}
+end
+
+"""
+    ModernNCALoader
+
+Training iterator yielding `((x, cand_x, cand_y, y), y)` batches. The corpus is
+moved to device once at construction; query rows follow a shuffled epoch
+permutation and candidates are resampled each step.
+
+# Fields
+- `full_x`, `full_y`: device corpus.
+- `batchsize`: query rows per step.
+- `n_cand`: candidate rows sampled per step (`0` when `batchsize == N`).
+- `rng`: RNG for sampling.
+- `dev`: device-transfer callable.
+"""
+struct ModernNCALoader{X,Y,R<:AbstractRNG,D}
+    full_x::X
+    full_y::Y
+    batchsize::Int
+    n_cand::Int
+    rng::R
+    dev::D
+end
+
+function ModernNCAConfig(;
+    d_embedding::Int=128, n_blocks::Int=2, d_block::Int=256,
+    dropout::Real=0.1, temperature::Real=1.0, sample_rate::Real=0.8,
+    max_candidates::Int=8192, eps::Real=1f-8,
+)
+    return ModernNCAConfig(d_embedding, n_blocks, d_block,
+        Float32(dropout), Float32(temperature), Float32(sample_rate),
+        max_candidates, Float32(eps))
+end
+
+"""
+    _backbone(cfg, d_in, embedding_layer)
+
+Build the ModernNCA encoder: embedding, linear, then n_blocks times (BN, Dense(relu), Dropout, Dense), then BN.
+"""
+function _backbone(cfg::ModernNCAConfig, d_in::Int, embedding_layer)
+    emb = isnothing(embedding_layer) ? WrappedFunction(identity) : embedding_layer
+    layers = Any[emb, Dense(d_in => cfg.d_embedding)]
+    for _ in 1:cfg.n_blocks
+        push!(layers, BatchNorm(cfg.d_embedding))
+        push!(layers, Dense(cfg.d_embedding => cfg.d_block, relu))
+        cfg.dropout > 0 && push!(layers, Dropout(cfg.dropout))
+        push!(layers, Dense(cfg.d_block => cfg.d_embedding))
+    end
+    cfg.n_blocks > 0 && push!(layers, BatchNorm(cfg.d_embedding))
+    return Chain(layers...)
+end
+
+function (cfg::ModernNCAConfig)(; nfeats, outsize,
+                                 loss_type::Type{<:LossType}=MSE,
+                                 embedding_layer=nothing)
+    return ModernNCAModel(_backbone(cfg, nfeats, embedding_layer), cfg, Int(outsize), loss_type)
+end
+
+Base.length(l::ModernNCALoader) = fld(size(l.full_x, 2), l.batchsize)
+
+"""
+    _pairwise_dist(q, k, ϵ) -> Matrix
+
+`(num_keys, batch)` Euclidean distance matrix between `q` `(d, batch)` and `k` `(d, num_keys)`.
+`ϵ` is added under `sqrt` for numerical stability.
+"""
+function _pairwise_dist(q::AbstractMatrix, k::AbstractMatrix, ϵ::Float32)
+    q2 = sum(abs2, q; dims=1)  # (1, batch)
+    k2 = sum(abs2, k; dims=1)  # (1, num_keys)
+    kq = k' * q                 # (num_keys, batch)
+    return sqrt.(max.(0f0, k2' .+ q2 .- 2f0 .* kq) .+ ϵ)
+end
+
+_diag_inf(i, j, d) = ifelse(i == j, typemax(typeof(d)), d)
+
+"""
+    _mask_diag(d) -> Matrix
+
+Set the diagonal of distance matrix `d` to `typemax`, collapsing self-attention
+weights to zero during training.
+"""
+_mask_diag(d::AbstractMatrix) =
+    _diag_inf.(reshape(1:size(d, 1), :, 1), reshape(1:size(d, 2), 1, :), d)
+
+_to_output(::Type{<:Union{MSE,MAE}}, α, cy, _) = reshape(α' * cy, 1, :)
+
+function _to_output(::Type{<:LogLoss}, α, cy, _)
+    p = clamp.(reshape(α' * cy, 1, :), 1f-6, 1f0 - 1f-6)
+    return log.(p ./ (1f0 .- p))
+end
+
+function _to_output(::Type{<:MLogLoss}, α, cy, outsize::Int)
+    oh = ((k, c) -> ifelse(k == c, 1f0, 0f0)).(
+        reshape(UInt32(1):UInt32(outsize), :, 1), reshape(cy, 1, :))
+    return log.(clamp.(oh * α, 1f-7, Inf32))
+end
+
+"""
+    _nca_logits(m, zq, zk, cy; mask_self=false) -> Matrix
+
+Softmax attention over corpus keys `zk` for query embeddings `zq`.
+When `mask_self=true`, diagonal entries are masked (training path).
+"""
+function _nca_logits(m::ModernNCAModel, zq, zk, cy; mask_self::Bool=false)
+    d = _pairwise_dist(zq, zk, m.cfg.eps) ./ max(m.cfg.temperature, m.cfg.eps)
+    mask_self && (d = _mask_diag(d))
+    α = softmax(-d; dims=1)
+    return _to_output(m.loss_type, α, cy, m.outsize)
+end
+
+"""
+    _encode_corpus(m, cx, ps, st; chunk=2048) -> Matrix
+
+Encode raw corpus `cx` `(d_in, N)` in chunks of `chunk` rows to bound peak
+memory and stay within BatchNorm shape limits. Returns `zk` `(d_embedding, N)`.
+"""
+function _encode_corpus(m::ModernNCAModel, cx, ps, st; chunk::Int=2048)
+    n = size(cx, 2)
+    if n <= chunk
+        return m.backbone(cx, ps, st)[1]
+    end
+    z0, _ = m.backbone(cx[:, 1:chunk], ps, st)
+    zk = similar(z0, size(z0, 1), n)
+    zk[:, 1:chunk] .= z0
+    s = chunk + 1
+    while s <= n
+        e = min(s + chunk - 1, n)
+        zk[:, s:e] .= m.backbone(cx[:, s:e], ps, st)[1]
+        s = e + 1
+    end
+    return zk
+end
+
+"""
+    (m::ModernNCAModel)((x, cand_x, cand_y, y), ps, st)
+
+Training forward: encodes queries and candidates separately, matching the
+reference ModernNCA BatchNorm behavior. The encoded queries are prepended to
+the candidate embeddings before NCA attention, then self-neighbors are masked.
+"""
+function (m::ModernNCAModel)((x, cand_x, cand_y, y)::Tuple{Any,Any,Any,Any}, ps, st)
+    zq, st_bb = m.backbone(x, ps, st)
+    zc, st_bb = size(cand_x, 2) == 0 ?
+        (similar(zq, size(zq, 1), 0), st_bb) :
+        m.backbone(cand_x, ps, st_bb)
+    z_all = hcat(zq, zc)
+    cy = vcat(vec(y), vec(cand_y))
+    return _nca_logits(m, zq, z_all, cy; mask_self=true), st_bb
+end
+
+"""
+    (m::ModernNCAModel)((x, cx, cy), ps, st)
+
+Inference/eval forward: encodes queries and the raw corpus `cx` with the current
+parameters, then attends over the corpus. Encoding is done here (not pre-computed)
+so the corpus embeddings always match the current model during evaluation.
+"""
+function (m::ModernNCAModel)((x, cand_x, cand_y)::Tuple{Any,Any,Any}, ps, st)
+    zq, st_bb = m.backbone(x, ps, st)
+    zk = _encode_corpus(m, cand_x, ps, st_bb)
+    return _nca_logits(m, zq, zk, cand_y), st_bb
+end
+
+"""
+    Base.iterate(l::ModernNCALoader, state=nothing)
+
+`state = (perm, start)` carries the epoch permutation across calls, so query
+batches are random without depending on dataframe row order.
+
+Candidates are sampled from the complement of the query batch within `perm`.
+The index-skip mapping avoids allocating that complement: for `j` in
+`1:(n-batchsize)`, use `perm[j]` before the batch window and
+`perm[j + batchsize]` after it. `n_cand == 0` only when `batchsize == n`; then
+the forward pass keys on the batch itself with diagonal masking.
+"""
+function Base.iterate(l::ModernNCALoader, state=nothing)
+    n = size(l.full_x, 2)
+    perm, start = state === nothing ? (randperm(l.rng, n), 1) : state
+    stop = start + l.batchsize - 1
+    stop > n && return nothing
+
+    batch_idx = perm[start:stop]
+
+    bidx = l.dev(batch_idx)
+    x = l.full_x[:, bidx]
+    y = l.full_y[bidx]
+
+    if l.n_cand > 0
+        m = n - l.batchsize
+        js = rand(l.rng, 1:m, l.n_cand)
+        cidx = l.dev(perm[@. ifelse(js < start, js, js + l.batchsize)])
+        cand_x = l.full_x[:, cidx]
+        cand_y = l.full_y[cidx]
+    else
+        cand_x = similar(l.full_x, size(l.full_x, 1), 0)
+        cand_y = similar(l.full_y, 0)
+    end
+    return ((x, cand_x, cand_y, y), y), (perm, stop + 1)
+end
+
+"""
+    build_corpus(df, feature_names, target_name, loss_type, scalers)
+
+Return `(full_x, full_y)`: the corpus feature matrix `(d_in, N)` as `Float32`
+and the encoded target vector, ready for `ModernNCALoader`.
+"""
+function build_corpus(df::AbstractDataFrame, feature_names, target_name,
+                      loss_type::Type{<:LossType}, scalers)
+    full_x = permutedims(Matrix{Float32}(select(df, collect(feature_names))))
+    full_y = _encode_targets(df, target_name, loss_type, scalers)
+    return full_x, full_y
+end
+
+"""
+    _encode_targets(df, target_name, loss_type, scalers)
+
+Encode targets for each loss type:
+- `MLogLoss`: 1-based `UInt32` class codes.
+- `LogLoss`: `Float32` in `{0, 1}`.
+- `MSE / MAE`: `Float32`, standardised when `scalers` is provided.
+
+# Arguments
+- `df`: source data frame.
+- `target_name`: target column.
+- `loss_type`: target encoding dispatch.
+- `scalers`: optional target scaler `(mu, sigma)`.
+"""
+_encode_targets(df, target_name, ::Type{<:MLogLoss}, _) =
+    UInt32.(CategoricalArrays.levelcode.(df[!, target_name]))
+
+function _encode_targets(df, target_name, ::Type{<:LogLoss}, _)
+    col = df[!, target_name]
+    eltype(col) <: CategoricalValue || return Float32.(col)
+    levels = CategoricalArrays.levels(col)
+    length(levels) == 2 || error("For `loss=:logloss`, target must have exactly 2 classes.")
+    return Float32.(CategoricalArrays.levelcode.(col) .- 1)
+end
+
+function _encode_targets(df, target_name, ::Type{<:Union{MSE,MAE,GaussianMLE}}, scalers)
+    y = Float32.(df[!, target_name])
+    isnothing(scalers) || (y .= (y .- scalers.mu) ./ scalers.sigma)
+    return y
+end
+
+"""
+    Models.build_chain(cfg::ModernNCAConfig, embed_chain; outsize, d_in, loss_type, kw...)
+
+Pass the embedding layer into the backbone instead of wrapping it in an outer `Chain`.
+"""
+Models.build_chain(cfg::ModernNCAConfig, embed_chain;
+        outsize, d_in, loss_type, kw...) =
+    cfg(; nfeats=d_in, outsize, loss_type, embedding_layer=embed_chain)
+
+"""
+    Models.train_dataloader(cfg::ModernNCAConfig, ...)
+
+Build `ModernNCALoader` and stash the raw corpus on `m.info[:nca_ref]`.
+`n_cand = floor(sample_rate × (N − batchsize))`, capped by
+`cfg.max_candidates` when positive. `sample_rate ≥ 1` uses the full complement
+before the cap; `batchsize == N` gives `n_cand = 0`.
+
+# Arguments
+- `cfg`: ModernNCA config.
+- `m`: fitted model wrapper.
+- `df`: training data frame.
+- `feature_names`, `target_name`: data columns.
+- `loss_type`, `scalers`: target encoding metadata.
+- `batchsize`, `dev`, `rng`: loader settings.
+"""
+function Models.train_dataloader(cfg::ModernNCAConfig, m::NeuroTabModel, ::Any, df;
+        feature_names, target_name, loss_type, scalers, batchsize, dev, rng, kw...)
+    cx, cy = build_corpus(df, feature_names, target_name, loss_type, scalers)
+    m.info[:nca_ref] = (cx=cx, cy=cy)
+    n = size(cx, 2)
+    batchsize = min(batchsize, n)
+    pool = n - batchsize
+    raw_n_cand = pool == 0 ? 0 :
+        cfg.sample_rate >= 1f0 ? pool :
+        max(Int(floor(cfg.sample_rate * pool)), 1)
+    n_cand = cfg.max_candidates > 0 ? min(raw_n_cand, cfg.max_candidates) : raw_n_cand
+    return ModernNCALoader(dev(cx), dev(cy), batchsize, n_cand, rng, dev)
+end
+
+"""
+    Models.infer_dataloader(m::ModernNCAModel, ...)
+
+Wrap each inference batch as `(x, cx, cy)` with the raw training corpus.
+The corpus is encoded inside the forward pass using the current parameters.
+"""
+function Models.infer_dataloader(::ModernNCAModel, info, data, dev, ps=nothing, st=nothing)
+    ref = info[:nca_ref]
+    cx, cy = dev(ref.cx), dev(ref.cy)
+    return Iterators.map(x -> (x, cx, cy), data)
+end
+
+"""
+    Models.eval_dataloader(m::ModernNCAModel, ...)
+
+Wrap each eval batch as `((x, cx, cy), rest...)` with the raw training corpus.
+The corpus is encoded inside the forward pass so embeddings always use the
+current model parameters rather than a stale pre-encoded version.
+"""
+function Models.eval_dataloader(::ModernNCAModel, info, data, dev, ps=nothing, st=nothing)
+    ref = info[:nca_ref]
+    cx, cy = dev(ref.cx), dev(ref.cy)
+    return Iterators.map(d -> ((d[1], cx, cy), Base.tail(d)...), data)
+end
+
+end # module
diff --git a/src/models/embeddings/batchnorm.jl b/src/models/embeddings/batchnorm.jl
index 464215c..538a96a 100644
--- a/src/models/embeddings/batchnorm.jl
+++ b/src/models/embeddings/batchnorm.jl
@@ -1,14 +1,17 @@
 """
-    BatchNormEmbeddings(nfeats)
+    _BatchNormEmbeddings(; nfeats)
 
 Feature-wise `BatchNorm` on `(nfeats, batch)` input. Output shape matches input.
+
+# Arguments
+- `nfeats::Int`: Number of input features.
 """
-struct BatchNormEmbeddings{L} <: LuxCore.AbstractLuxWrapperLayer{:layer}
+struct _BatchNormEmbeddings{L} <: LuxCore.AbstractLuxWrapperLayer{:layer}
     layer::L
 end
 
-BatchNormEmbeddings(nfeats::Int) = BatchNormEmbeddings(BatchNorm(nfeats))
+_BatchNormEmbeddings(; nfeats::Int) = _BatchNormEmbeddings(BatchNorm(nfeats))
 
-function (l::BatchNormEmbeddings)(x::AbstractMatrix, ps, st)
-    return l.layer(x, ps, st)
-end
+(l::_BatchNormEmbeddings)(x::AbstractMatrix, ps, st) = l.layer(x, ps, st)
+
+LuxCore.outputsize(l::_BatchNormEmbeddings, x, ::AbstractRNG) = (size(x, 1),)
\ No newline at end of file
diff --git a/src/models/embeddings/compute_bins.jl b/src/models/embeddings/compute_bins.jl
index 157d030..48e69a5 100644
--- a/src/models/embeddings/compute_bins.jl
+++ b/src/models/embeddings/compute_bins.jl
@@ -3,7 +3,7 @@
 
 Quantile-based bin edges for piecewise-linear embeddings.
 
-`X` must have shape `(n_samples, n_features)` — the transpose of model input `(n_features, batch)`.
+`X` must have shape `(n_samples, n_features)`,the transpose of model input `(n_features, batch)`.
 
 # Arguments
 - `X::AbstractMatrix`: Training data `(n_samples, n_features)`.
@@ -37,7 +37,13 @@ function compute_bins(X::AbstractMatrix; bins::Union{Int,Vector{Int}}=48)
         quantile_probs = range(0f0, 1f0, length=n_bins_vec[j] + 1)
         edges = Float32[quantile(col_buf, p; sorted=true) for p in quantile_probs]
         unique!(edges)
-        @assert length(edges) >= 2 "Feature $j has fewer than 2 unique bin edges"
+        if length(edges) < 2
+            # Constant features collapse all quantile edges to one value.
+            # Give the encoder a tiny nonzero-width bin instead of failing.
+            v = edges[1]
+            delta = max(abs(v) * 1f-3, 1f-3)
+            edges = Float32[v - delta, v + delta]
+        end
         bins[j] = edges
     end
     return bins
diff --git a/src/models/embeddings/config.jl b/src/models/embeddings/config.jl
index cc78ef6..bb38a8b 100644
--- a/src/models/embeddings/config.jl
+++ b/src/models/embeddings/config.jl
@@ -7,77 +7,372 @@ const act_dict = Dict(
     :tanhshrink => tanhshrink,
 )
 
+# Supertype for embedding specs the learner can hold (`EmbeddingLayer`, numerical embeddings).
+abstract type AbstractEmbedding end
+
+# Column-wise numerical-feature embeddings; also the type of an `EmbeddingLayer`'s `num`.
+abstract type AbstractNumericalEmbedding <: AbstractEmbedding end
+
+# Time-column embeddings.
+abstract type AbstractTemporalEmbedding end
+
+"""
+    IdentityEmbedding()
+
+No-op numerical embedding that passes features through unchanged; the default `num` for
+an [`EmbeddingLayer`](@ref).
+"""
+struct IdentityEmbedding <: AbstractNumericalEmbedding end
+
 """
-    EmbeddingConfig(; kwargs...)
+    LinearEmbeddings(; d_embedding=16, activation=:relu)
 
-Configuration for numerical feature embeddings.
-Pass as `embedding_config` to `NeuroTabRegressor` / `NeuroTabClassifier`.
+Per-feature linear embedding: each feature is projected to `d_embedding` by its own affine map, followed by `activation`.
 
 # Arguments
-- `embedding_type::Symbol`: `:periodic`, `:linear`, `:piecewise`, or `:batchnorm` (default `:periodic`).
-- `d_embedding::Int`: Embedding dimension per feature (default `16`; forced to `1` for `:batchnorm`).
-- `activation::Symbol`: `:identity`, `:relu`, `:tanh`, etc. (default `:relu`, or `:identity` for `:piecewise`).
-- `bins::Union{Int, Vector{Int}}`: Bin count for `:piecewise` (default `32`).
-- `frequencies::Int`: Sinusoidal components for `:periodic` (default `32`).
-- `frequencies_init_scale::Float32`: σ for periodic frequency init (default `0.01f0`).
-"""
-struct EmbeddingConfig
-    embedding_type::Symbol
+- `d_embedding::Int`: Output dimension per feature (default `16`).
+- `activation::Symbol`: Activation applied after the projection (default `:relu`).
+"""
+struct LinearEmbeddings <: AbstractNumericalEmbedding
     d_embedding::Int
     activation::Symbol
-    bins::Union{Int,Vector{Int}}
+end
+function LinearEmbeddings(; d_embedding::Int=16, activation::Symbol=:relu)
+    d_embedding > 0 || throw(ArgumentError("d_embedding must be > 0, got $d_embedding"))
+    LinearEmbeddings(d_embedding, activation)
+end
+
+"""
+    PeriodicEmbeddings(; d_embedding=16, frequencies=32, frequencies_init_scale=0.01f0,
+                       activation=:relu, lite=false)
+
+Per-feature periodic embedding: each feature is expanded with learned sine/cosine components, then projected to `d_embedding` and passed through `activation`.
+
+# Arguments
+- `d_embedding::Int`: Output dimension per feature (default `16`).
+- `frequencies::Int`: Number of sinusoidal components per feature (default `32`).
+- `frequencies_init_scale::Float32`: Std. dev. initializing the frequencies (default `0.01f0`).
+- `activation::Symbol`: Activation applied after the projection (default `:relu`).
+- `lite::Bool`: Share the projection across features to cut parameters (default `false`).
+"""
+struct PeriodicEmbeddings <: AbstractNumericalEmbedding
+    d_embedding::Int
     frequencies::Int
     frequencies_init_scale::Float32
+    activation::Symbol
+    lite::Bool
 end
+function PeriodicEmbeddings(; d_embedding::Int=16, frequencies::Int=32,
+                            frequencies_init_scale::Real=0.01f0, activation::Symbol=:relu,
+                            lite::Bool=false)
+    d_embedding > 0 || throw(ArgumentError("d_embedding must be > 0, got $d_embedding"))
+    frequencies > 0 || throw(ArgumentError("frequencies must be > 0, got $frequencies"))
+    PeriodicEmbeddings(d_embedding, frequencies, Float32(frequencies_init_scale),
+                       activation, lite)
+end
+
+"""
+    PiecewiseLinearEmbeddings(; d_embedding=16, bins=32, activation=:identity, version=:B)
 
-function EmbeddingConfig(;
-    embedding_type=:periodic,
+Per-feature piecewise-linear embedding against bin edges computed from the training data, then projected to `d_embedding`. The bin edges are derived at fit time, so this embedding requires `x_train`.
+
+# Arguments
+- `d_embedding::Int`: Output dimension per feature (default `16`).
+- `bins::Union{Int,Vector{Int}}`: Number of bins, or per-feature bin counts (default `32`).
+- `activation::Symbol`: Activation applied after the projection (default `:identity`).
+- `version::Symbol`: Encoding variant, `:A` or `:B` (default `:B`).
+"""
+struct PiecewiseLinearEmbeddings <: AbstractNumericalEmbedding
+    d_embedding::Int
+    bins::Union{Int,Vector{Int}}
+    activation::Symbol
+    version::Symbol
+end
+function PiecewiseLinearEmbeddings(; d_embedding::Int=16, bins::Union{Int,Vector{Int}}=32,
+                                   activation::Symbol=:identity, version::Symbol=:B)
+    d_embedding > 0 || throw(ArgumentError("d_embedding must be > 0, got $d_embedding"))
+    version in (:A, :B) ||
+        throw(ArgumentError("version must be :A or :B, got :$version"))
+    PiecewiseLinearEmbeddings(d_embedding, bins, activation, version)
+end
+
+"""
+    BatchNormEmbeddings()
+
+Batch-normalize the raw features without expanding their dimension (one output per feature).
+"""
+struct BatchNormEmbeddings <: AbstractNumericalEmbedding end
+
+"""
+    TemporalEmbeddings(; index, order=[4, 1, 7, 0], periods=_DEFAULT_TEMPORAL_PERIODS,
+                       trend=true, d_embedding=16)
+
+Fourier embedding of a single time column: the column at `index` is expanded into multi-scale sine/cosine features at `periods`, projected to `d_embedding`, and optionally augmented with a linear trend. `order` and `periods` align per band  `order[i]` is the harmonic count for `periods[i]`.
+
+# Arguments
+- `index::Int`: Required. 1-based position of the time column in `feature_names`.
+- `order::Vector{Int}`: Harmonics per period; nonnegative with at least one positive entry (default `[4, 1, 7, 0]`).
+- `periods::Vector{Float32}`: Base periods in column units (default `_DEFAULT_TEMPORAL_PERIODS`).
+- `trend::Bool`: Append a linear trend term to the embedding (default `true`).
+- `d_embedding::Int`: Projection dimension of the periodic features (default `16`).
+"""
+struct TemporalEmbeddings <: AbstractTemporalEmbedding
+    index::Int
+    order::Vector{Int}
+    periods::Vector{Float32}
+    trend::Bool
+    d_embedding::Int
+end
+function TemporalEmbeddings(;
+    index::Int,
+    order::AbstractVector{<:Integer}=Int[4, 1, 7, 0],
+    periods::AbstractVector{<:Real}=_DEFAULT_TEMPORAL_PERIODS,
+    trend::Bool=true,
     d_embedding::Int=16,
-    activation=nothing,
-    bins::Union{Int,Vector{Int}}=32,
-    frequencies::Int=32,
-    frequencies_init_scale::Float32=0.01f0,
+)
+    index >= 1 ||
+        throw(ArgumentError("index must be >= 1, got $index"))
+    length(order) == length(periods) ||
+        throw(ArgumentError("length(order)=$(length(order)) must equal length(periods)=$(length(periods))"))
+    all(>=(0), order) && any(>(0), order) ||
+        throw(ArgumentError("order must be non-negative with at least one positive entry"))
+    d_embedding > 0 ||
+        throw(ArgumentError("d_embedding must be > 0, got $d_embedding"))
+    TemporalEmbeddings(index, Vector{Int}(order), Vector{Float32}(periods),
+                       trend, d_embedding)
+end
+
+"""
+    EmbeddingLayer(; num=IdentityEmbedding(), temp=nothing)
+    EmbeddingLayer(num::AbstractNumericalEmbedding; temp=nothing)
+    EmbeddingLayer(d::AbstractDict)
+
+Numerical embedding over the non-time columns, optionally combined with a temporal
+embedding on the column at `temp.index`. `num` always exists (defaults to
+[`IdentityEmbedding`](@ref), raw pass-through); `nothing` normalizes to it. When `temp`
+is set the branches are concatenated features-first, temporal-last.
+
+The `Dict` form is for the hyperparameter harness: `:embedding_type` (`:linear`,
+`:periodic`, `:piecewise`, `:batchnorm`, `:identity`) selects the numerical type, an
+optional `:temporal` Dict gives [`TemporalEmbeddings`](@ref) kwargs, and unknown or
+`nothing` keys are ignored (so a superset Dict works; missing `:embedding_type` → identity).
+
+```julia
+EmbeddingLayer(PeriodicEmbeddings(; d_embedding=24))
+EmbeddingLayer(; temp=TemporalEmbeddings(; index=1))            # temporal only
+EmbeddingLayer(Dict(:embedding_type => :periodic, :temporal => Dict(:index => 1)))
+```
+"""
+struct EmbeddingLayer{
+    N<:AbstractNumericalEmbedding,
+    T<:Union{Nothing,AbstractTemporalEmbedding},
+} <: AbstractEmbedding
+    num::N
+    temp::T
+end
+_as_num(::Nothing) = IdentityEmbedding()
+_as_num(n::AbstractNumericalEmbedding) = n
+EmbeddingLayer(; num=IdentityEmbedding(), temp=nothing) = EmbeddingLayer(_as_num(num), temp)
+EmbeddingLayer(num::AbstractNumericalEmbedding; temp=nothing) = EmbeddingLayer(num, temp)
+
+const _NUM_EMBEDDING_TYPES = Dict{Symbol,Type}(
+    :identity  => IdentityEmbedding,
+    :linear    => LinearEmbeddings,
+    :periodic  => PeriodicEmbeddings,
+    :piecewise => PiecewiseLinearEmbeddings,
+    :batchnorm => BatchNormEmbeddings,
 )
 
-    embedding_type = Symbol(embedding_type)
-    if isnothing(activation)
-        activation = embedding_type == :piecewise ? :identity : :relu
+# Used by the `_num_from_dict` methods below. Filters a (possibly superset) config
+# Dict down to the keyword arguments one embedding constructor accepts; absent and
+# `nothing`-valued keys are skipped so the constructor's own defaults apply.
+function _pick(d, keys)
+    ps = Pair{Symbol,Any}[]
+    for k in keys
+        haskey(d, k) && d[k] !== nothing && push!(ps, k => d[k])
     end
-    activation = Symbol(activation)
+    return ps
+end
 
-    if embedding_type == :batchnorm
-        d_embedding = 1
+_num_from_dict(::Type{IdentityEmbedding}, d) = IdentityEmbedding()
+_num_from_dict(::Type{LinearEmbeddings}, d) =
+    LinearEmbeddings(; _pick(d, (:d_embedding, :activation))...)
+_num_from_dict(::Type{PeriodicEmbeddings}, d) =
+    PeriodicEmbeddings(; _pick(d, (:d_embedding, :frequencies, :frequencies_init_scale, :activation, :lite))...)
+_num_from_dict(::Type{PiecewiseLinearEmbeddings}, d) =
+    PiecewiseLinearEmbeddings(; _pick(d, (:d_embedding, :bins, :activation, :version))...)
+_num_from_dict(::Type{BatchNormEmbeddings}, d) = BatchNormEmbeddings()
+
+function EmbeddingLayer(d::AbstractDict)
+    d = Dict{Symbol,Any}(Symbol(k) => v for (k, v) in d)
+    num = if haskey(d, :embedding_type) && d[:embedding_type] !== nothing
+        etype = Symbol(d[:embedding_type])
+        T = get(_NUM_EMBEDDING_TYPES, etype) do
+            throw(ArgumentError("unknown :embedding_type $(repr(etype)); valid: $(collect(keys(_NUM_EMBEDDING_TYPES)))"))
+        end
+        _num_from_dict(T, d)
+    else
+        IdentityEmbedding()
     end
+    temp = if haskey(d, :temporal)
+        td = Dict{Symbol,Any}(Symbol(k) => v for (k, v) in d[:temporal])
+        TemporalEmbeddings(; td...)
+    else
+        nothing
+    end
+    return EmbeddingLayer(num, temp)
+end
+
+# Output width of a temporal embedding: d_embedding, +1 when trend is appended.
+temporal_out_dim(t::TemporalEmbeddings) = t.d_embedding + (t.trend ? 1 : 0)
+
+# Per-feature output width of a numerical embedding. Identity and BatchNorm keep each
+# feature at width 1; the expanding embeddings emit `d_embedding` columns per feature.
+_num_d_embedding(::IdentityEmbedding) = 1
+_num_d_embedding(::BatchNormEmbeddings) = 1
+_num_d_embedding(s::Union{LinearEmbeddings,PeriodicEmbeddings,PiecewiseLinearEmbeddings}) = s.d_embedding
+
+"""
+    needs_x_train(spec) -> Bool
 
-    return EmbeddingConfig(embedding_type, d_embedding, activation, bins, frequencies, frequencies_init_scale)
+Whether building `spec` requires the training matrix.
+
+# Arguments
+- `spec::AbstractEmbedding`: The embedding spec to inspect.
+
+# Returns
+`true` when `spec` needs `x_train` (e.g. piecewise-linear bin edges or temporal normalization statistics), `false` otherwise.
+"""
+needs_x_train(::Nothing) = false
+needs_x_train(::IdentityEmbedding) = false
+needs_x_train(::AbstractNumericalEmbedding) = false
+needs_x_train(::PiecewiseLinearEmbeddings) = true
+needs_x_train(::AbstractTemporalEmbedding) = true
+needs_x_train(e::EmbeddingLayer) = needs_x_train(e.num) || needs_x_train(e.temp)
+
+# Numerical embeddings: each builds its own chain emitting a flat (width, batch)
+# output. Expanding types append their own FlattenLayer; BatchNorm is already 2D and
+# Identity is a no-op. The builder never inspects the type to decide how to flatten;
+# adding a new numerical embedding means adding one method here.
+# Row (feature-column) selector used by the temporal `Parallel` branches. Named (not a
+# closure) so the resulting `WrappedFunction` is concretely typed via `Base.Fix2`.
+_select_rows(x::AbstractMatrix, rows) = x[rows, :]
+
+_build_num(::IdentityEmbedding; nfeats::Int, x_train=nothing) = NoOpLayer()
+_build_num(s::LinearEmbeddings; nfeats::Int, x_train=nothing) =
+    Chain(_LinearEmbeddings(; nfeats, d_embedding=s.d_embedding, activation=act_dict[s.activation]),
+        FlattenLayer())
+_build_num(s::PeriodicEmbeddings; nfeats::Int, x_train=nothing) =
+    Chain(_PeriodicEmbeddings(; nfeats, d_embedding=s.d_embedding,
+        frequencies=s.frequencies, frequencies_init_scale=s.frequencies_init_scale,
+        activation=act_dict[s.activation], lite=s.lite), FlattenLayer())
+function _build_num(s::PiecewiseLinearEmbeddings; nfeats::Int, x_train=nothing)
+    x_train === nothing &&
+        error("PiecewiseLinearEmbeddings requires x_train to compute bin edges")
+    Chain(_PiecewiseLinearEmbeddings(; bins=compute_bins(x_train; bins=s.bins),
+        d_embedding=s.d_embedding, activation=act_dict[s.activation], version=s.version),
+        FlattenLayer())
+end
+_build_num(s::BatchNormEmbeddings; nfeats::Int, x_train=nothing) = _BatchNormEmbeddings(; nfeats)
+
+function _build_temp(t::TemporalEmbeddings, x_col)
+    x_col === nothing &&
+        error("TemporalEmbeddings requires x_train for t_mean/t_std")
+    t_mean = Float32(mean(x_col))
+    t_std = length(x_col) > 1 ? max(Float32(std(x_col)), 1f-6) : 1f0
+    _TemporalEmbeddings(t_mean, t_std, t.order, t.trend, t.d_embedding; periods=t.periods)
 end
 
 """
-    (config::EmbeddingConfig)(; nfeats, x_train=nothing)
+    build_embedding_chain(spec, nfeats; x_train=nothing)
 
-Build `Chain(embedding, FlattenLayer())` from `config`.
+Build the embedding chain for `spec`. A numerical-only spec embeds every column; with a
+temporal branch the input is routed through a `Lux.Parallel`, applying `num` to the non-time
+columns and `temp` to the time column, concatenated features-first then temporal.
 
 # Arguments
+- `spec::AbstractEmbedding`: The embedding spec to build.
 - `nfeats::Int`: Number of input features.
-- `x_train`: Training matrix `(n_samples, n_features)`; required for `:piecewise`.
-"""
-function (config::EmbeddingConfig)(; nfeats::Int, x_train=nothing)
-    emb = if config.embedding_type == :periodic
-        PeriodicEmbeddings(nfeats, config.d_embedding;
-            frequencies=config.frequencies,
-            frequencies_init_scale=config.frequencies_init_scale,
-            activation=act_dict[config.activation])
-    elseif config.embedding_type == :linear
-        LinearEmbeddings(nfeats, config.d_embedding; activation=act_dict[config.activation])
-    elseif config.embedding_type == :piecewise
-        @assert x_train !== nothing "Piecewise embeddings require `x_train` to compute bin edges."
-        bins = compute_bins(x_train; bins=config.bins)
-        @assert length(bins) == nfeats "Expected $nfeats bin vectors, got $(length(bins))"
-        PiecewiseLinearEmbeddings(bins, config.d_embedding; activation=act_dict[config.activation])
-    elseif config.embedding_type == :batchnorm
-        BatchNormEmbeddings(nfeats)
-    else
-        error("Unsupported embedding type: $(config.embedding_type)")
-    end
-    return Chain(emb, FlattenLayer())
+- `x_train`: Training matrix `(n_samples, nfeats)`, required when `needs_x_train(spec)` (default `nothing`).
+
+# Returns
+A `Lux` layer emitting a flat `(width, batch)` output; recover the width with [`embedding_width`](@ref).
+"""
+build_embedding_chain(s::AbstractNumericalEmbedding, nfeats::Int; x_train=nothing) =
+    _build_num(s; nfeats, x_train)
+function build_embedding_chain(e::EmbeddingLayer, nfeats::Int; x_train=nothing)
+    # numerical branch only
+    isnothing(e.temp) && return _build_num(e.num; nfeats, x_train)
+
+    # numerical + temporal: route the non-time columns through `num` and the time
+    # column through `temp`, concatenating features-first / temporal-last. When `num`
+    # is IdentityEmbedding this is the temporal-only case (features pass through).
+    idx = e.temp.index
+    1 <= idx <= nfeats ||
+        throw(ArgumentError("temporal index=$idx out of range for nfeats=$nfeats"))
+    keep = setdiff(1:nfeats, idx)
+    core = _build_num(e.num; nfeats=nfeats - 1, x_train=x_train === nothing ? nothing : x_train[:, keep])
+    temp = _build_temp(e.temp, x_train === nothing ? nothing : @view x_train[:, idx])
+    return Parallel(vcat,
+        Chain(WrappedFunction(Base.Fix2(_select_rows, keep)), core),
+        Chain(WrappedFunction(Base.Fix2(_select_rows, idx:idx)), temp))
+end
+
+"""
+    embedding_width(layer, x, rng) -> Int
+
+Output width of `layer`, measured by a single forward pass on the probe `x`. A forward pass
+is robust to the layer's internal structure, unlike analytic `outputsize` which cannot see
+through a `vcat` connection.
+
+# Arguments
+- `layer`: The embedding layer to measure.
+- `x::AbstractMatrix`: Probe input of shape `(nfeats, batch)` (`init` passes `(nfeats, 2)`).
+- `rng::AbstractRNG`: RNG used for parameter setup.
+
+# Returns
+The flattened output width as an `Int`.
+"""
+function embedding_width(layer, x, rng::AbstractRNG)
+    ps, st = LuxCore.setup(rng, layer)
+    st = LuxCore.testmode(st)
+    y, _ = layer(x, ps, st)
+    return size(y, 1)
 end
+
+"""
+    has_real_embedding(spec) -> Bool
+
+Whether `spec` applies a non-identity embedding.
+
+# Arguments
+- `spec::AbstractEmbedding`: The embedding spec to inspect.
+
+# Returns
+`true` when `spec` has a non-identity `num` or a temporal branch; `false` for a bare `IdentityEmbedding` or an empty `EmbeddingLayer`.
+"""
+has_real_embedding(::IdentityEmbedding) = false
+has_real_embedding(::AbstractNumericalEmbedding) = true
+has_real_embedding(e::EmbeddingLayer) = !(e.num isa IdentityEmbedding) || e.temp !== nothing
+
+"""
+    per_feature_widths(spec, nfeats) -> Vector{Int}
+
+Output width contributed by each input feature, summing to the total embedding width. Used by
+TabM's grouped scaling init to keep a feature's columns together. Ordering matches the concat
+layout: non-time features first (numerical per-feature width), temporal column last.
+
+# Arguments
+- `spec::AbstractEmbedding`: The embedding spec.
+- `nfeats::Int`: Number of input features.
+
+# Returns
+A `Vector{Int}` of per-feature widths summing to the total embedding width.
+"""
+per_feature_widths(::IdentityEmbedding, nfeats::Int) = fill(1, nfeats)
+per_feature_widths(s::AbstractNumericalEmbedding, nfeats::Int) = fill(_num_d_embedding(s), nfeats)
+function per_feature_widths(e::EmbeddingLayer, nfeats::Int)
+    d_emb = _num_d_embedding(e.num)
+    isnothing(e.temp) && return fill(d_emb, nfeats)
+    return vcat(fill(d_emb, nfeats - 1), Int[temporal_out_dim(e.temp)])
+end
\ No newline at end of file
diff --git a/src/models/embeddings/embeddings.jl b/src/models/embeddings/embeddings.jl
index 918f231..801860f 100644
--- a/src/models/embeddings/embeddings.jl
+++ b/src/models/embeddings/embeddings.jl
@@ -1,16 +1,16 @@
 module Embeddings
 
 using Lux
-using Lux: BatchNorm, Chain, Dense, FlattenLayer
 using LuxCore
+using Random
 using NNlib
-using Random: AbstractRNG, rand, randn
-using Statistics: quantile
+import Statistics: mean, quantile, std
 
-export NLinear, LinearEmbeddings
-export Periodic, PeriodicEmbeddings
-export PiecewiseLinearEncoding, PiecewiseLinearEmbeddings
-export compute_bins, EmbeddingConfig
+export AbstractNumericalEmbedding, AbstractTemporalEmbedding, AbstractEmbedding
+export LinearEmbeddings, PeriodicEmbeddings, PiecewiseLinearEmbeddings
+export BatchNormEmbeddings, TemporalEmbeddings, IdentityEmbedding
+export EmbeddingLayer, build_embedding_chain, needs_x_train, temporal_out_dim
+export per_feature_widths, has_real_embedding, embedding_width
 
 include("compute_bins.jl")
 include("nlinear.jl")
@@ -18,6 +18,7 @@ include("linear.jl")
 include("periodic.jl")
 include("piecewise_linear.jl")
 include("batchnorm.jl")
+include("temporal.jl")
 include("config.jl")
 
-end
+end
\ No newline at end of file
diff --git a/src/models/embeddings/linear.jl b/src/models/embeddings/linear.jl
index da7d589..32de0cc 100644
--- a/src/models/embeddings/linear.jl
+++ b/src/models/embeddings/linear.jl
@@ -1,24 +1,27 @@
 """
-    LinearEmbeddings(nfeats, d_embedding; activation=relu)
+    _LinearEmbeddings(; nfeats, d_embedding, activation=relu)
 
-Per-feature affine map `activation(w * x + b)`. Output shape `(d_embedding, nfeats, batch)`.
+Embeds each continuous feature via a learned affine transformation followed by
+an activation: `activation(w_j * x_j + b_j)`.
+Produces a `(d_embedding, nfeats, batch)` tensor.
 
 # Arguments
 - `nfeats::Int`: Number of input features.
 - `d_embedding::Int`: Embedding dimension per feature.
-- `activation`: Element-wise activation (default `relu`).
+- `activation`: Activation function applied element-wise (default `relu`).
+  E.g. `relu`, `tanh`, `identity`.
 """
-struct LinearEmbeddings{F} <: LuxCore.AbstractLuxLayer
+struct _LinearEmbeddings{F} <: LuxCore.AbstractLuxLayer
     nfeats::Int
     d_embedding::Int
     activation::F
 end
 
-function LinearEmbeddings(nfeats::Int, d_embedding::Int; activation=relu)
-    return LinearEmbeddings(nfeats, d_embedding, activation)
+function _LinearEmbeddings(; nfeats::Int, d_embedding::Int, activation=NNlib.relu)
+    return _LinearEmbeddings(nfeats, d_embedding, activation)
 end
 
-function LuxCore.initialparameters(rng::AbstractRNG, l::LinearEmbeddings)
+function LuxCore.initialparameters(rng::AbstractRNG, l::_LinearEmbeddings)
     limit = Float32(l.d_embedding)^(-0.5f0)
     weight = reshape((rand(rng, Float32, l.d_embedding, l.nfeats) .* 2f0 .* limit) .- limit,
         l.d_embedding, l.nfeats, 1)
@@ -27,9 +30,11 @@ function LuxCore.initialparameters(rng::AbstractRNG, l::LinearEmbeddings)
     return (weight=weight, bias=bias)
 end
 
-LuxCore.initialstates(::AbstractRNG, ::LinearEmbeddings) = (;)
+LuxCore.initialstates(::AbstractRNG, ::_LinearEmbeddings) = (;)
 
-function (l::LinearEmbeddings)(x::AbstractMatrix, ps, st)
+function (l::_LinearEmbeddings)(x::AbstractMatrix, ps, st)
     x_r = reshape(x, 1, size(x, 1), size(x, 2))
     return l.activation.(muladd.(ps.weight, x_r, ps.bias)), st
 end
+
+LuxCore.outputsize(l::_LinearEmbeddings, x, ::AbstractRNG) = (l.d_embedding, size(x, 1))
\ No newline at end of file
diff --git a/src/models/embeddings/periodic.jl b/src/models/embeddings/periodic.jl
index a96a144..90c16d4 100644
--- a/src/models/embeddings/periodic.jl
+++ b/src/models/embeddings/periodic.jl
@@ -1,12 +1,13 @@
 """
     Periodic(nfeats, n_frequencies, sigma)
 
-Sinusoidal encoding `[cos(2π w x), sin(2π w x)]`. Output shape `(2 * n_frequencies, nfeats, batch)`.
+Maps each feature to `2 * n_frequencies` sinusoidal components
+`[cos(2π w x), sin(2π w x)]`. Output shape `(2 * n_frequencies, nfeats, batch)`.
 
 # Arguments
 - `nfeats::Int`: Number of input features.
-- `n_frequencies::Int`: Frequency components per feature.
-- `sigma::Float32`: Std-dev for frequency weight init (clamped to ±3σ).
+- `n_frequencies::Int`: Number of frequency components per feature.
+- `sigma::Float32`: Std-dev for the frequency weight initialization (clamped to ±3σ).
 """
 struct Periodic <: LuxCore.AbstractLuxLayer
     nfeats::Int
@@ -30,32 +31,34 @@ function (l::Periodic)(x::AbstractMatrix, ps, st)
 end
 
 """
-    PeriodicEmbeddings(nfeats, d_embedding=24; frequencies=48,
-                       frequencies_init_scale=0.01f0, activation=relu, lite=false)
+    _PeriodicEmbeddings(; nfeats, d_embedding=24, frequencies=48,
+                        frequencies_init_scale=0.01f0, activation=relu, lite=false)
 
-`Periodic` followed by per-feature linear projection and activation.
+Periodic sinusoidal encoding followed by a learned linear projection: applies
+`Periodic`, then `NLinear` (or a shared `Dense` if `lite`), then `activation`.
 
 # Arguments
 - `nfeats::Int`: Number of input features.
-- `d_embedding::Int`: Output dimension per feature (default `24`).
-- `frequencies::Int`: Sinusoidal components per feature (default `48`).
-- `frequencies_init_scale::Float32`: σ for frequency init (default `0.01f0`).
-- `activation`: Post-projection activation (default `relu`).
-- `lite::Bool`: Use shared `Dense` instead of per-feature `NLinear` (default `false`).
+- `d_embedding::Int`: Output embedding dimension per feature (default `24`).
+- `frequencies::Int`: Sinusoidal frequency components per feature (default `48`).
+- `frequencies_init_scale::Float32`: σ for frequency weight init (default `0.01f0`).
+- `activation`: Activation applied after projection (default `relu`).
+- `lite::Bool`: Use a single shared `Dense` instead of per-feature `NLinear`
+  (default `false`). Requires a non-identity `activation`.
 """
-struct PeriodicEmbeddings{P,L,F} <: LuxCore.AbstractLuxContainerLayer{(:periodic, :linear)}
+struct _PeriodicEmbeddings{P,L,F} <: LuxCore.AbstractLuxContainerLayer{(:periodic, :linear)}
     periodic::P
     linear::L
     activation::F
     lite::Bool
 end
 
-function PeriodicEmbeddings(
+function _PeriodicEmbeddings(;
     nfeats::Int,
-    d_embedding::Int=24;
+    d_embedding::Int=24,
     frequencies::Int=48,
     frequencies_init_scale::Float32=0.01f0,
-    activation=relu,
+    activation=NNlib.relu,
     lite::Bool=false,
 )
     if lite && activation === identity
@@ -67,10 +70,10 @@ function PeriodicEmbeddings(
     else
         NLinear(nfeats, 2 * frequencies, d_embedding)
     end
-    return PeriodicEmbeddings(periodic, linear, activation, lite)
+    return _PeriodicEmbeddings(periodic, linear, activation, lite)
 end
 
-function (m::PeriodicEmbeddings)(x::AbstractMatrix, ps, st)
+function (m::_PeriodicEmbeddings)(x::AbstractMatrix, ps, st)
     h, st_p = m.periodic(x, ps.periodic, st.periodic)
 
     h_lin, st_l = if m.lite
@@ -82,5 +85,10 @@ function (m::PeriodicEmbeddings)(x::AbstractMatrix, ps, st)
         m.linear(h, ps.linear, st.linear)
     end
 
-    return m.activation.(h_lin), (periodic=st_p, linear=st_l)
+    h_lin = m.activation.(h_lin)
+
+    return h_lin, (periodic=st_p, linear=st_l)
 end
+
+LuxCore.outputsize(m::_PeriodicEmbeddings, x, ::AbstractRNG) =
+    m.lite ? error("lite _PeriodicEmbeddings outputsize undefined") : (m.linear.out_features, size(x, 1))
\ No newline at end of file
diff --git a/src/models/embeddings/piecewise_linear.jl b/src/models/embeddings/piecewise_linear.jl
index 223a7a4..bfa332f 100644
--- a/src/models/embeddings/piecewise_linear.jl
+++ b/src/models/embeddings/piecewise_linear.jl
@@ -1,7 +1,7 @@
 """
     PiecewiseLinearEncoding(bins)
 
-Non-trainable piecewise-linear encoding from precomputed bin edges.
+Non-trainable piecewise-linear encoding using precomputed bin edges.
 Output shape `(max_n_bins, nfeats, batch)`.
 
 # Arguments
@@ -36,6 +36,9 @@ function LuxCore.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding)
         b = -bin_edges[1:end-1] ./ bin_width
         nb = length(bin_edges) - 1
 
+        # Place the last bin's weight/bias at the end row;
+        # remaining bins fill rows 1:nb-1. Unused rows stay zero
+        # and are clamped to [0, 1] harmlessly.
         weight[end, i] = w[end]
         bias[end, i] = b[end]
         if nb > 1
@@ -44,6 +47,7 @@ function LuxCore.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding)
         end
     end
 
+    # Pre-reshape to 3D to avoid per-call reshape in forward pass
     return (
         weight=reshape(weight, M, N, 1),
         bias=reshape(bias, M, N, 1),
@@ -51,25 +55,27 @@ function LuxCore.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding)
 end
 
 function (l::PiecewiseLinearEncoding)(x::AbstractMatrix, ps, st)
+    # For each feature j and bin b: h[b,j,:] = clamp(w[b,j]*x[j,:] + bias[b,j], 0, 1)
     x_r = reshape(x, 1, size(x, 1), size(x, 2))
     h = clamp.(muladd.(st.weight, x_r, st.bias), 0f0, 1f0)
     return h, st
 end
 
 """
-    PiecewiseLinearEmbeddings(bins, d_embedding; activation=identity, version=:B)
+    _PiecewiseLinearEmbeddings(; bins, d_embedding, activation=identity, version=:B)
 
-Learnable projection on top of [`PiecewiseLinearEncoding`](@ref).
-Version `:A`: encoding → `NLinear` (with bias).
-Version `:B`: encoding → zero-init `NLinear` + [`LinearEmbeddings`](@ref) residual.
+Learnable embeddings on top of `PiecewiseLinearEncoding`.
+Version `:A`: PLE -> NLinear (with bias).
+Version `:B`: PLE -> NLinear (zero-init, no bias) + per-feature linear residual.
+Output shape `(d_embedding, nfeats, batch)`.
 
 # Arguments
-- `bins::Vector{<:AbstractVector}`: Bin edges from [`compute_bins`](@ref).
-- `d_embedding::Int`: Output dimension per feature.
-- `activation`: Post-projection activation (default `identity`).
+- `bins::Vector{<:AbstractVector}`: Bin edges per feature from [`compute_bins`](@ref).
+- `d_embedding::Int`: Embedding dimension per feature.
+- `activation`: Activation function applied after projection (default `identity`). E.g. `relu`, `tanh`.
 - `version::Symbol`: `:A` or `:B` (default `:B`).
 """
-struct PiecewiseLinearEmbeddings{L0,I,L,F} <: LuxCore.AbstractLuxContainerLayer{(:linear0, :encoding, :linear)}
+struct _PiecewiseLinearEmbeddings{L0,I,L,F} <: LuxCore.AbstractLuxContainerLayer{(:linear0, :encoding, :linear)}
     linear0::L0
     encoding::I
     linear::L
@@ -77,9 +83,9 @@ struct PiecewiseLinearEmbeddings{L0,I,L,F} <: LuxCore.AbstractLuxContainerLayer{
     version::Symbol
 end
 
-function PiecewiseLinearEmbeddings(
+function _PiecewiseLinearEmbeddings(;
     bins::Vector{<:AbstractVector},
-    d_embedding::Int;
+    d_embedding::Int,
     activation=identity,
     version::Symbol=:B,
 )
@@ -88,13 +94,14 @@ function PiecewiseLinearEmbeddings(
     max_n_bins = maximum(length(b) - 1 for b in bins)
 
     encoding = PiecewiseLinearEncoding(bins)
-    linear0 = (version == :B) ? LinearEmbeddings(nfeats, d_embedding; activation=identity) : nothing
+    # Residual path uses raw affine output (no activation)
+    linear0 = (version == :B) ? _LinearEmbeddings(; nfeats, d_embedding, activation=identity) : nothing
     linear = NLinear(nfeats, max_n_bins, d_embedding; bias=(version == :A))
 
-    return PiecewiseLinearEmbeddings(linear0, encoding, linear, activation, version)
+    return _PiecewiseLinearEmbeddings(linear0, encoding, linear, activation, version)
 end
 
-function LuxCore.initialparameters(rng::AbstractRNG, m::PiecewiseLinearEmbeddings)
+function LuxCore.initialparameters(rng::AbstractRNG, m::_PiecewiseLinearEmbeddings)
     ps_l0 = m.linear0 === nothing ? nothing : LuxCore.initialparameters(rng, m.linear0)
     ps_enc = LuxCore.initialparameters(rng, m.encoding)
 
@@ -108,7 +115,7 @@ function LuxCore.initialparameters(rng::AbstractRNG, m::PiecewiseLinearEmbedding
     return (linear0=ps_l0, encoding=ps_enc, linear=ps_lin)
 end
 
-function (m::PiecewiseLinearEmbeddings)(x::AbstractMatrix, ps, st)
+function (m::_PiecewiseLinearEmbeddings)(x::AbstractMatrix, ps, st)
     val_linear0 = nothing
     st_l0 = st.linear0
 
@@ -124,3 +131,5 @@ function (m::PiecewiseLinearEmbeddings)(x::AbstractMatrix, ps, st)
 
     return h_final, (linear0=st_l0, encoding=st_enc, linear=st_lin)
 end
+
+LuxCore.outputsize(m::_PiecewiseLinearEmbeddings, x, ::AbstractRNG) = (m.linear.out_features, size(x, 1))
\ No newline at end of file
diff --git a/src/models/embeddings/temporal.jl b/src/models/embeddings/temporal.jl
new file mode 100644
index 0000000..fd37433
--- /dev/null
+++ b/src/models/embeddings/temporal.jl
@@ -0,0 +1,77 @@
+"""
+Default Fourier base periods for `TemporalEmbeddings`, in seconds:
+year (`365.25 * 86_400`), 30-day month, week, day. Assumes the time column is
+encoded as a POSIX-style seconds-since-epoch float.
+"""
+const _DEFAULT_TEMPORAL_PERIODS = Float32[31_557_600, 2_629_800, 604_800, 86_400]
+
+"""Fixed Fourier features `sin(ωᵢ t), cos(ωᵢ t)` per harmonic; `ωᵢ = 2πk/T`."""
+struct TemporalPeriodic <: LuxCore.AbstractLuxLayer
+    omega::Vector{Float32}
+end
+
+LuxCore.initialparameters(::AbstractRNG, ::TemporalPeriodic) = (;)
+
+# `omega` is fixed (no gradient), but it lives in `st` rather than as a struct
+# field so that `dev(st)` transfers it to the GPU alongside the learnable
+# state. Reshaped to a column so the broadcast against `(1, batch)` x is unambiguous.
+LuxCore.initialstates(::AbstractRNG, l::TemporalPeriodic) =
+    (omega = Matrix{Float32}(reshape(l.omega, length(l.omega), 1)),)
+
+function (l::TemporalPeriodic)(x::AbstractMatrix, ps, st)
+    z = st.omega .* x
+    return vcat(sin.(z), cos.(z)), st
+end
+
+"""Realized temporal embedding: Fourier features + dense projection (+ optional linear trend)."""
+struct _TemporalEmbeddings{P,D,trend} <: LuxCore.AbstractLuxContainerLayer{(:periodic, :dense)}
+    periodic::P
+    dense::D
+    t_mean::Float32
+    t_std::Float32
+end
+
+"""
+    _TemporalEmbeddings(t_mean, t_std, order, trend, d_embedding; periods)
+
+Realize a `_TemporalEmbeddings` layer from spec.
+
+`order` and `periods` are aligned per-band: for each `(o_i, p_i)`, the Fourier
+basis contributes harmonics `k = 1:o_i` with angular frequency `ω = 2πk/p_i`.
+All `ω` are concatenated into a single flat vector, so the resulting
+`TemporalPeriodic` outputs `2 * sum(order)` features (sin + cos per harmonic),
+which the `Dense` then projects to `d_embedding`.
+
+`trend` is lifted into the type parameter so the forward dispatches statically
+on whether the standardised raw time `(x - t_mean) / t_std` is appended.
+"""
+function _TemporalEmbeddings(
+    t_mean::Real, t_std::Real,
+    order::AbstractVector{<:Integer}, trend::Bool, d_embedding::Int;
+    periods::AbstractVector{<:Real}=_DEFAULT_TEMPORAL_PERIODS,
+)
+    @assert length(order) == length(periods) "length(order) must equal length(periods)"
+    @assert any(>(0), order) "`order` must contain at least one positive entry"
+    omega = Float32[2f0 * Float32(π) * Float32(k) / Float32(p)
+                    for (o, p) in zip(order, periods) for k in 1:o]
+    periodic = TemporalPeriodic(omega)
+    dense = Dense(2 * length(omega) => d_embedding, NNlib.relu)
+    return _TemporalEmbeddings{typeof(periodic),typeof(dense),trend}(
+        periodic, dense, Float32(t_mean), Float32(t_std),
+    )
+end
+
+function (m::_TemporalEmbeddings{P,D,true})(x::AbstractMatrix, ps, st) where {P,D}
+    h, st_p = m.periodic(x, ps.periodic, st.periodic)
+    out, st_d = m.dense(h, ps.dense, st.dense)
+    return vcat(out, (x .- m.t_mean) ./ m.t_std), (periodic=st_p, dense=st_d)
+end
+
+function (m::_TemporalEmbeddings{P,D,false})(x::AbstractMatrix, ps, st) where {P,D}
+    h, st_p = m.periodic(x, ps.periodic, st.periodic)
+    out, st_d = m.dense(h, ps.dense, st.dense)
+    return out, (periodic=st_p, dense=st_d)
+end
+
+LuxCore.outputsize(m::_TemporalEmbeddings{P,D,true},  x, ::AbstractRNG) where {P,D} = (m.dense.out_dims + 1,)
+LuxCore.outputsize(m::_TemporalEmbeddings{P,D,false}, x, ::AbstractRNG) where {P,D} = (m.dense.out_dims,)
\ No newline at end of file
diff --git a/src/models/models.jl b/src/models/models.jl
index 3e8e3d1..73cbf5e 100644
--- a/src/models/models.jl
+++ b/src/models/models.jl
@@ -1,8 +1,11 @@
 module Models
 
 export NeuroTabModel, Architecture
-export Embeddings
-export NeuroTreeConfig, MLPConfig, ResNetConfig, TabMConfig, MOETreeConfig
+export Embeddings, EmbeddingLayer
+export LinearEmbeddings, PeriodicEmbeddings, PiecewiseLinearEmbeddings
+export BatchNormEmbeddings, TemporalEmbeddings, IdentityEmbedding
+export AbstractNumericalEmbedding, AbstractTemporalEmbedding, AbstractEmbedding
+export NeuroTreeConfig, MLPConfig, ResNetConfig, TabMConfig, MOETreeConfig, ModernNCAConfig
 
 using ..Losses
 using Lux: Chain
@@ -32,16 +35,60 @@ function get_activation(act::Symbol)
     return activation_dict[act]
 end
 
+"""
+    train_dataloader(arch, m, default, df; kwargs...)
+
+Per-architecture hook: return the dataloader `fit` should use. Default returns
+the tabular `default` unchanged; retrieval-style archs override to substitute
+their own iterator (e.g. with a candidate corpus attached).
+
+Overrides may write into `m.info` (e.g. to stash a corpus for inference).
+"""
+train_dataloader(::Architecture, ::Any, data, ::Any; kw...) = data
+
+"""
+    build_chain(arch, embed_chain; nfeats, outsize, d_in, kw...)
+
+Per-architecture hook for assembling the Lux chain that `fit` will train.
+The embedding chain always exists (identity is `WrappedFunction(identity)`);
+`d_in` is supplied by the caller. Default wires `embed_chain` in front of the
+architecture's backbone via `Chain(embed_chain, arch(...))`.
+"""
+build_chain(arch::Architecture, embed_chain; nfeats, outsize, d_in, kw...) =
+    Chain(embed_chain, arch(; nfeats=d_in, outsize))
+
+"""
+    infer_dataloader(chain, info, data, dev, ps, st)
+
+Per-architecture hook: return the per-batch iterator `infer` should use.
+Default returns `data` unchanged; retrieval-style archs override to wrap each
+batch with extra inputs (e.g. a candidate corpus). Dispatched on
+`typeof(m.chain)` so arch modules can override on their concrete model type.
+"""
+infer_dataloader(::Any, ::Any, data, ::Any, ::Any, ::Any) = data
+
+"""
+    eval_dataloader(chain, info, data, dev, ps, st)
+
+Per-architecture hook: return the per-batch iterator the eval `CallBack` should
+use. Default returns `data` unchanged; retrieval-style archs override to wrap
+each batch's `x` with extra inputs (e.g. a candidate corpus) so that the eval
+forward pass matches the training/inference signature. Dispatched on
+`typeof(m.chain)`.
+"""
+eval_dataloader(::Any, ::Any, data, ::Any, ::Any, ::Any) = data
+
 """
     NeuroTabModel
 
 The object containing the model and associated metadata.
 
-## Fields
+# Fields
 
-- `loss_type`: the loss function type used during training (e.g. `MSE`, `LogLoss`, `MLogLoss`, `GaussianMLE`)
-- `chain`: the underlying `Lux.Chain` neural network
-- `info`: a `Dict{Symbol,Any}` storing metadata such as `:feature_names`, `:target_levels`, `:device`, `logger`, as well as fitted parameters (`ps`) and state (`st`).
+- `loss_type`: the loss type used in training (`MSE`, `LogLoss`, `MLogLoss`, `GaussianMLE`).
+- `chain`: the underlying `Lux.Chain` neural network.
+- `info`: a `Dict{Symbol,Any}` of metadata such as `:feature_names`, `:target_levels`,
+  and `:device`, plus the fitted parameters (`ps`) and state (`st`).
 """
 struct NeuroTabModel{L<:LossType,C}
     loss_type::Type{L}
@@ -61,10 +108,19 @@ using .MOETrees
 include("TabM/TabM.jl")
 using .TabM
 
+function build_chain(arch::TabMConfig, embed_chain; nfeats, outsize, d_in, embedding_config, kw...)
+    d_features = per_feature_widths(embedding_config, nfeats)
+    scaling_init_override = has_real_embedding(embedding_config) ? :normal : nothing
+    return Chain(embed_chain, arch(; nfeats=d_in, outsize, d_features, scaling_init_override))
+end
+
 include("MLP/mlp.jl")
 using .MLP
 
 include("ResNet/resnet.jl")
 using .ResNet
 
+include("ModernNCA/modernnca.jl")
+using .ModernNCA
+
 end
\ No newline at end of file
diff --git a/test/core.jl b/test/core.jl
index f7ac49a..a0dc557 100644
--- a/test/core.jl
+++ b/test/core.jl
@@ -192,7 +192,7 @@ end
         embedding_config=Dict(:embedding_type => :batchnorm),
         nrounds=200,
         batchsize=32,
-        early_stopping_rounds=5,
+        early_stopping_rounds=10,
         lr=1e-2,
     )
 
diff --git a/test/embedding.jl b/test/embedding.jl
index 1e73a5e..e98413c 100644
--- a/test/embedding.jl
+++ b/test/embedding.jl
@@ -77,7 +77,7 @@ end
             early_stopping_rounds=20,
             embedding_config)
 
-        m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names, print_every_n=5)
+        m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names, print_every_n=5);
 
         ptrain = [argmax(x) for x in eachrow(m(dtrain))]
         peval = [argmax(x) for x in eachrow(m(deval))]