diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py
index d0c1cf2..d9b8e5c 100644
--- a/astrai/config/model_config.py
+++ b/astrai/config/model_config.py
@@ -20,6 +20,7 @@ class BaseModelConfig(BaseConfig):
     """Base config with ``model_type`` dispatch and file I/O."""
 
     model_type: Optional[str] = None
+    neftune_alpha: float = 0.0
 
 
 @dataclass
diff --git a/astrai/model/components/embedding.py b/astrai/model/components/embedding.py
index f8f4551..4d400fb 100644
--- a/astrai/model/components/embedding.py
+++ b/astrai/model/components/embedding.py
@@ -7,10 +7,13 @@ from torch import Tensor
 
 
 class Embedding(nn.Module):
-    def __init__(self, vocab_size: int, embedding_dim: int):
+    def __init__(self, vocab_size: int, embedding_dim: int, neftune_alpha: float = 0.0):
         super().__init__()
         self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
-        self.neftune_noise_alpha = 0.0
+        self.neftune_noise_alpha = neftune_alpha
+
+    def set_neftune_alpha(self, alpha: float):
+        self.neftune_noise_alpha = alpha
 
     def reset_parameters(self):
         nn.init.normal_(self.weight, mean=0.0, std=0.02)
diff --git a/astrai/model/encoder.py b/astrai/model/encoder.py
index 1ea1b2c..4b83652 100644
--- a/astrai/model/encoder.py
+++ b/astrai/model/encoder.py
@@ -23,7 +23,9 @@ class EmbeddingEncoder(AutoModel):
         self.rotary_embedding = RotaryEmbedding(
             rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
         )
-        self.embed_tokens = Embedding(config.vocab_size, config.dim)
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+        )
 
         self.layers = nn.ModuleList(
             [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py
index 44655f8..9810545 100644
--- a/astrai/model/transformer.py
+++ b/astrai/model/transformer.py
@@ -59,7 +59,9 @@ class AutoRegressiveLM(AutoModel):
         self.rotary_embedding = RotaryEmbedding(
             rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
         )
-        self.embed_tokens = Embedding(config.vocab_size, config.dim)
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+        )
 
         self.layers = nn.ModuleList(
             [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py
index 031ab51..c7c7877 100644
--- a/astrai/trainer/train_context.py
+++ b/astrai/trainer/train_context.py
@@ -63,7 +63,6 @@ class TrainContextBuilder:
 
         model = cfg.model_fn()
         model = model.to(device=device)
-        model.embed_tokens.neftune_noise_alpha = cfg.neftune_alpha
 
         model_config = {}
         if self._resume_dir:
diff --git a/scripts/tools/train.py b/scripts/tools/train.py
index e845596..999ff0a 100644
--- a/scripts/tools/train.py
+++ b/scripts/tools/train.py
@@ -309,6 +309,7 @@ def train(
     # Load config
     config_path = os.path.join(param_path, "config.json")
     config = AutoRegressiveLMConfig.from_file(config_path)
+    config.neftune_alpha = neftune_alpha
 
     if window_size is None:
         window_size = config.max_len