{"id":"f3718625-33a2-4dcb-8f00-d93deb274fa3","arxiv_id":"2501.08888","paper_version":2,"verdict":"CONDITIONAL","confidence":"HIGH","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":6,"one_line_summary":"A two-stage pretraining-finetuning method with partial initialization of prediction heads improves CATE estimation from combined observational and RCT data under hidden confounding.","lead":"This paper proposes a two-stage training method that first learns a representation from large observational data, then fine-tunes on a small randomized trial to correct hidden confounding. The key trick is initializing part of the prediction network from the first stage to avoid overfitting on the small trial.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"No ablation isolates the proposed partial initialization from the L_shift regularizer, so Table 1 cannot support the claim that TSPF avoids overfitting via initialization; the central mechanism is untested.","rationale":"The reader's weakest assumption concerns whether the finetuned heads initialized from biased observational outcomes provide a better starting point than random initialization, and whether L_shift prevents the RCT data from removing hidden confounding. This is exactly the untested premise I identify: the paper's central mechanism is the partial initialization plus the L_shift regularizer, but the experiments never isolate these components. A method can be competitive in Table 1 for reasons unrelated to the proposed initialization, such as the frozen representation or the small adapter, without validating the paper's stated contribution. I do not find a fatal internal inconsistency in the math: the block-zero initialization does make the initial predictions match the first-stage model, and the gradients are nonzero after the first step, so the construction is coherent. The main weakness is evidential, not logical. The validation-set conversion to a randomized distribution is a secondary concern because the true CATE is unchanged under this relabeling, though it weakens the claim that out-of-sample gains reflect the original data-generating process. The reader's conditional verdict is appropriate: the paper needs an ablation of the initialization strategy and the regularizer before the central overfitting claim can be accepted. My concern reinforces the reader's verdict rather than moving it, so the verdict should remain unchanged.","tokens_in":13582,"tokens_out":6098,"duration_ms":69442,"concrete_test":"Run a controlled ablation on IHDP and Jobs with the same two-stage pipeline but four variants: (a) proposed partial initialization, (b) random initialization of g0 and g1, (c) full initialization of g0 and g1 from h0 and h1, and (d) proposed initialization with lambda_4 = 0 (no L_shift), across RCT training fractions of 5%, 10%, and 20%. Report out-of-sample sqrt(PEHE) and the in-sample/out-of-sample gap. If variant (a) does not beat (b) and (c) by a nontrivial margin, or if (d) matches (a), the claimed overfitting mitigation is not attributable to the partial initialization.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The paper's central assertion is that the partial initialization strategy in the second stage 'can avoid overfitting' caused by small RCT data. The only evidence is Table 1, which compares complete methods. Nothing varies the initialization while holding the architecture, the frozen representation, the adapter, the mutual-information penalty, and L_shift fixed. The observed gains could come from the frozen representation plus a small adapter, from the L_shift proximity penalty, from the mutual-information term, or from the initialization itself; Table 1 cannot distinguish these. Moreover, if L_shift dominates, the prediction heads remain close to the first-stage observational predictors, so the RCT signal cannot remove hidden confounding; if the initialization is unnecessary, the headline claim is unsupported. Section 3.2 asserts that the initialization strategy is crucial but provides no ablation, and Section 4.2 converts the validation set to a randomized distribution, so out-of-sample performance is not evaluated under the original observational assignment. Thus the load-bearing condition—that starting from biased observational heads is better than random initialization and that the bias is removable under L_shift—is not established by the reported experiments.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes a two-stage pretraining-finetuning framework (TSPF) for estimating conditional average treatment effects when observational data are confounded by unmeasured variables and only a small randomized controlled trial (RCT) is available. In the first stage, a representation network and two outcome heads are trained on large observational data. In the second stage, the representation is frozen, a small adapter network is introduced, the representation and adapter outputs are concatenated, and the two outcome heads are finetuned on the RCT data with a partial initialization taken from the first-stage heads. The method also uses a CLUB mutual-information penalty between the two representations and an L_shift regularizer that anchors the finetuned heads to their initialized values. Experiments on two semi-synthetic datasets (IHDP and Jobs) report lower PEHE and ATE errors than a set of baselines, and the paper claims that the partial initialization mitigates overfitting from the small RCT sample.","tokens_in":13899,"tokens_out":6612,"duration_ms":69931,"significance":"If the central claim holds, the paper addresses a practical and important problem: combining biased large-scale observational data with small unconfounded RCT data for heterogeneous treatment effect estimation. The two-stage design is intuitive, and the partial initialization idea is a reasonable mechanism to combat overfitting in small RCT finetuning. The experimental comparison on two standard benchmarks is a useful empirical contribution, and the paper explicitly discusses the most related residual-correction baseline. However, the paper's own claim about the mechanism is not isolated by any ablation, and the synthetic experiments are confined to a linear and additive data-generating process, which undercuts the paper's stated generality. The significance of the work therefore depends on additional experiments that the authors have not reported.","major_comments":[{"comment":"The manuscript's central claim is that the partial initialization strategy in the second stage 'can avoid overfitting' caused by small RCT data. Table 1 only compares complete methods (TSPF versus baselines); no experiment varies the initialization scheme while holding the frozen representation φ, the adapter φ_U, the CLUB mutual-information penalty, and the L_shift regularization fixed. Without such an ablation, the observed gains cannot be attributed to the partial initialization; they could come from the frozen representation with a small adapter, from the mutual-information term, or from L_shift. Please add ablation studies that at least compare (a) partial initialization versus random initialization, (b) L_shift on versus off, and (c) combinations, with the same architecture and data.","section":"Section 4.4, Table 1"},{"comment":"The synthetic data-generating process is linear and additive: μ_0 = W_3·X + W_4·U, μ_1 = W_5·X + W_6·U + 4, with Gaussian noise and a logistic propensity. The contributions and conclusion claim that 'the proposed framework does not rely on the linear and additive generation assumptions' and that the method has the advantage of 'not restricting the data-generating process.' This claim is not supported by the experiments, which only exercise this linear setting. Please add nonlinear synthetic experiments (e.g., neural or polynomial outcome and propensity functions) or revise the claim to reflect the tested scope.","section":"Section 4.2"},{"comment":"The L_shift penalty is defined relative to the first-stage initialization θ_g^0, which comes from a model trained on confounded observational data. If λ_4 is large, the finetuned heads remain close to that biased initialization and the RCT signal may be unable to remove the hidden confounding. The paper does not report sensitivity to λ_4, nor the final distance ||θ_g − θ_g^0|| after training, so it is unclear whether the regularizer is actually preventing overfitting or inadvertently preserving confounding bias. Please add a sensitivity analysis over λ_4 and, if possible, an analysis of the parameter displacement.","section":"Section 3.2, Regularization"},{"comment":"The validation set is transformed by replacing the original treatment and factual outcome with an RCT-like assignment before evaluation. If the validation set is used for early stopping or hyperparameter selection, model selection is performed under an unconfounded assignment rather than the original observational distribution. The paper should clarify how the validation set was used and report out-of-sample performance under the original observational assignment, or explain why the replacement does not affect model selection.","section":"Section 4.2"}],"minor_comments":[{"comment":"The row label 'TSFP (ours)' should be 'TSPF (ours)'.","section":"Table 1"},{"comment":"The activation function σ is used throughout the network definitions but is never specified; please state which activation function is used.","section":"Section 3.2.2"},{"comment":"Kallus et al. [23] is discussed as the closest residual-correction baseline but is not included in the experiments; please add it or justify its omission.","section":"Section 4.3.1"},{"comment":"Hyperparameter tuning is described only as 'tune the scale parameters from 1e−5 to 0.1'; provide the full search grid and selected values for reproducibility.","section":"Section 4.3.3"},{"comment":"The phrase 'significantly outperforms' is used without any statistical significance tests; report paired tests or confidence intervals across the repeated runs.","section":"Section 4.4"},{"comment":"The hidden confounder dimension is fixed at c = 30 and the RCT sample size is not varied; a sensitivity analysis over c and RCT sample size would strengthen the claim about overfitting mitigation.","section":"Section 4.2"}],"recommendation":"major_revision","confidential_remarks":"The manuscript is a workshop-style paper whose central contribution is an initialization scheme. The main risk is that the reported gains in Table 1 are not attributed to that scheme by any ablation, so the paper currently overclaims its mechanism. If the authors add the ablations and nonlinear simulations, the contribution could be a solid empirical study; as written, the load-bearing evidence is missing."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Quick take: the paper's empirical wins are real, but the claimed mechanism—partial initialization—is never isolated, and the evaluation protocol has a couple of red flags. If you cite it, cite the configuration, not the overfitting story.\n\nWhat's actually new: a two-stage pretraining-finetuning setup for CATE with hidden confounding. Stage one trains a representation and two heads on large OBS data. Stage two freezes that representation, adds a smaller adapter representation, concatenates them, widens the heads, copies the pretrained weights into the shared part, and zero-initializes the new rows. On top of that it adds a CLUB mutual-information penalty between the two representations and an L2 penalty on how far the heads move from their initialization. That combination doesn't appear in the cited Kallus et al. or CorNet work, and the empirical results on IHDP and Jobs are strong—TSPF is clearly better than the baselines on most metrics, and the gap over CorNet is large.\n\nWhat's soft. First, the central claim is about the partial initialization mitigating overfitting from small RCT data, but no ablation isolates it. Table 1 compares complete frameworks; the gains could come from the frozen representation plus adapter, from the CLUB term, from the L_shift regularizer, or from the initialization. The paper says the initialization is crucial but gives no experiment that varies it. Second, the data-generating process in Section 4.2 is linear: every term is a dot product W·x or W·u plus Gaussian noise. The paper claims generality to nonlinear mechanisms, but the experiments don't exercise any. Third, the validation set is converted to a randomized distribution by reassigning treatment with a fair coin and swapping in the counterfactual outcome. That means out-of-sample results are on a distribution with no confounding, not the observational distribution the method is supposed to correct. Fourth, in Table 1 the CorNet row is identical to the DRCFR row on every single metric. Either they didn't actually run CorNet and copied, or there's a reporting error. That alone makes the experimental table unusable as a comparison to the key two-stage baseline.\n\nBottom line: the idea is worth discussing and the configuration is new, but the evidence for the proposed mechanism is not there yet. A serious referee should ask for an ablation of the initialization, a fix for the CorNet/DRCFR duplicate, and an evaluation that respects the observational validation distribution. I'd send it to review—the core question is answerable with a few experiments—but not accept it as-is.","headline":"Real empirical wins, but the paper never isolates the partial initialization it credits; a duplicated baseline row and a randomized validation set further undercut the central claim.","tokens_in":14359,"tokens_out":3122,"would_cite":false,"duration_ms":32578,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":["62D20","68T07"],"pacs":[],"model":"deepseek-v4-flash","headline":"Two-stage pretraining-finetuning with partial head initialization avoids small-RCT overfitting and delivers competitive CATE estimates under hidden confounding.","keywords":["CATE estimation","hidden confounding","observational and RCT data fusion","two-stage pretraining-finetuning","partial parameter initialization","representation learning","mutual information regularization","overfitting mitigation"],"falsifier":"A direct ablation that resets the shared blocks of $g_0$ and $g_1$ to random values while keeping the adapter, the shift regularizer, the loss weights, and the optimizer identical would settle the claim: if the random-initialized version matches or beats TSPF's out-of-sample $\\sqrt{\\epsilon_{\\mathrm{PEHE}}}$ on IHDP or Jobs, then the partial initialization is not the source of the overfitting reduction.","tokens_in":13436,"feed_emoji":"🎯","tokens_out":7599,"duration_ms":75300,"temperature":0.7,"pith_summary":"Estimating conditional average treatment effects (CATE) from observational data is threatened by unmeasured confounders, while randomized controlled trials are unconfounded but often too small to train neural predictors. This paper tries to establish that a two-stage pretraining-finetuning strategy solves that tension: first fit a covariate representation and outcome heads on large observational data, then finetune on small RCT data using an adapter representation and prediction heads partially initialized from the first stage. The proposed TSPF framework is claimed to avoid overfitting from the small RCT sample and to produce the most competitive CATE estimates among the compared methods on the IHDP and Jobs benchmarks. If correct, this would make it practical to combine abundant biased data with scarce unbiased data for causal effect estimation without imposing linear or additive data-generation assumptions.","feed_headline":"Warm-started heads cut small-RCT overfitting in treatment-effect models","feed_subtitle":"Partial init lets tiny RCT data fix hidden confounding rather than memorize noise.","key_machinery":"The load-bearing mechanism is the partial initialization of the second-stage heads, supported by the representation adapter and the shift regularizer. Because the new head $g_t$ has wider layers than the pretrained $h_t$ to ingest the concatenated representation $[Z; Z_U]$, its parameters are partitioned into a shared block (same shape as $h_t$) and new blocks; shared weights and biases are copied from $h_t$ and every new block is zeroed. At initialization, the augmented branch therefore contributes nothing and the model's predictions coincide with the first-stage model, so finetuning on RCT data starts from the pretrained solution rather than from noise. The $\\mathcal{L}_{\\mathrm{shift}}$ penalty then constrains $g_0,g_1$ to stay close to that starting point, and the CLUB mutual-information term prevents the adapter representation $Z_U$ from simply copying $Z$. Together these pieces are intended to let the RCT data remove the hidden-confounding bias while preserving the covariate structure learned from the large observational dataset.","core_discovery":"The central claim is that the partial initialization of the second-stage prediction heads is what lets the model learn an unbiased CATE from small RCT data without overfitting. In the first stage, the representation module and heads $h_0,h_1$ are trained on observational data, which contain hidden confounding and therefore give biased outcomes. In the second stage, the first-stage representation is frozen; an adapter module produces an augmented representation $Z_U$, the two are concatenated, and new heads $g_0,g_1$ are finetuned on RCT data. The shared block of each $g_t$ weight matrix is initialized to $h_t$'s weights, the new rows and columns are set to zero, and the bias vectors are initialized similarly, so the finetuned model starts by reproducing the first-stage predictions. A shift regularizer keeps the finetuned parameters near this initialization, and a contrastive mutual-information bound keeps $Z_U$ from duplicating $Z$. The paper reports that this design outperforms the compared baselines, markedly beating the related two-stage method CorNet on both datasets.","pith_inferences":["Beyond the paper, a direct ablation holding the adapter, loss weights, and optimizer fixed while resetting the shared head blocks to random values would isolate whether the partial initialization itself, rather than the architecture or the shift regularizer, is responsible for the reported gains.","Beyond the paper, the shift regularizer may still bias the final estimate toward the observational solution, so an open scaling question is how much RCT data is needed to overcome that bias as confounding strength grows.","Beyond the paper, the same two-stage recipe should transfer to binary or survival outcomes, since the first-stage prediction loss already supports cross-entropy and the adapter width is tunable."],"forward_implications":["If the claim holds, CATE models can exploit large biased observational datasets plus tiny RCT samples without needing untestable assumptions such as instrumental variables or front-door adjustment.","Finetuning only the initialized heads and a small adapter means the effective number of parameters trained on RCT data is small, which is precisely the mechanism claimed to curb overfitting.","The adapter's width and depth can be shrunk or grown with the RCT sample size, giving a practical lever for overfitting control.","The method's advantage over CorNet suggests that learning a residual representation, rather than only estimating a residual function, is a productive direction for data fusion in causal inference."],"supporting_citations":[{"why":"Supplies the representation-balancing architecture and the Jobs benchmark that the first stage builds on.","marker":"[37]"},{"why":"Defines the residual-correction approach whose linear and additive assumptions the paper's adapter representation relaxes.","marker":"[23]"},{"why":"CorNet is the two-stage OBS+RCT baseline that TSPF is directly compared against and outperforms.","marker":"[13]"},{"why":"Provides the CLUB upper bound used to keep the adapter representation from duplicating the first-stage representation.","marker":"[5]"},{"why":"Provides the IHDP semi-synthetic benchmark used for the main experiments.","marker":"[16]"}],"fun_headline_variants":["Partial head init stops small-RCT overfitting in CATE","Warm-started heads tame tiny RCT data in causal models","Two-stage pretrain-finetune with partial init beats overfit","Start from observational, adapt heads to fix confounding"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The claim rests on the assumption that prediction heads pretrained on biased observational outcomes are a better starting point for learning the true CATE from RCT data than random initialization, and that the shift penalty keeping them close to that biased starting point will not simply preserve the confounding bias.","fun_headline_variants_meta":{"raw":{"variants":["Partial head init stops small-RCT overfitting in CATE","Warm-started heads tame tiny RCT data in causal models","Two-stage pretrain-finetune with partial init beats overfit","Start from observational, adapt heads to fix confounding"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.00066,"raw_usage":{"total_tokens":3033,"prompt_tokens":976,"completion_tokens":2057,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":592,"completion_tokens_details":{"reasoning_tokens":1988}},"tokens_in":592,"tokens_out":2057,"duration_ms":17268,"temperature":1.0,"reasoning_tokens":1988,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-10T20:15:09.658917+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"A direct ablation that resets the shared blocks of $g_0$ and $g_1$ to random values while keeping the adapter, the shift regularizer, the loss weights, and the optimizer identical would settle the claim: if the random-initialized version matches or beats TSPF's out-of-sample $\\sqrt{\\epsilon_{\\mathrm{PEHE}}}$ on IHDP or Jobs, then the partial initialization is not the source of the overfitting reduction.","supporting_citations":[{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Supplies the representation-balancing architecture and the Jobs benchmark that the first stage builds on."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Defines the residual-correction approach whose linear and additive assumptions the paper's adapter representation relaxes."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Provides the IHDP semi-synthetic benchmark used for the main experiments."}],"review_version":1}