diff --git a/source/MainComponent.cpp b/source/MainComponent.cpp
index 0cc4cef..e26478d 100644
--- a/source/MainComponent.cpp
+++ b/source/MainComponent.cpp
@@ -507,7 +507,7 @@ void MainComponent::buttonClicked (Button* buttonThatWasClicked)
         //[UserButtonCode_ShakeButton] -- add your button handler code here..
     	m_pLayer->weightsInit(weightInitLabel->getText().getFloatValue());
 		m_pLayer->redrawWeights();
-    	m_pLayer->redrawReconstruction();
+    	m_pLayer->redrawReconstruction(m_stack->trainingData(m_pLayer).row(m_trainingIndex));
         //[/UserButtonCode_ShakeButton]
     }
     else if (buttonThatWasClicked == createButton)
@@ -642,7 +642,7 @@ void MainComponent::sliderValueChanged (Slider* sliderThatWasMoved)
     {
         //[UserSliderCode_numGibbsSlider] -- add your slider handling code here..
     	m_pLayer->params().numGibbs = sliderThatWasMoved->getValue();
-		m_pLayer->redrawReconstruction();
+		m_pLayer->redrawReconstruction(m_stack->trainingData().row(m_trainingIndex));
         //[/UserSliderCode_numGibbsSlider]
     }
     else if (sliderThatWasMoved == m_progressBarSlider)
@@ -842,7 +842,7 @@ bool MainComponent::onProgress(Rbm *pRbm, const Rbm::Status &status)
 {
 	RbmComponent *pComp = static_cast<RbmComponent*>(pRbm);
 	pComp->upPass(pComp->DrawTraining->getData());
-	pComp->redrawReconstruction();
+	pComp->redrawReconstruction(pComp->DrawTraining->getData());
 	pComp->redrawWeights();
 
 	m_progressBarSlider->setValue(status.progress);
diff --git a/source/RbmComponent.cpp b/source/RbmComponent.cpp
index 38aa1c0..5af22ef 100644
--- a/source/RbmComponent.cpp
+++ b/source/RbmComponent.cpp
@@ -41,7 +41,7 @@ RbmComponent::RbmComponent (const std::string &name, size_t id, size_t numVisibl
     addAndMakeVisible (DrawHidden = new DrawComponent (numHidden, 1));
 	DrawHidden->setListener(this);
 	redrawWeights();
-	redrawReconstruction();
+//	redrawReconstruction();
 
     setSize (430, 130);
 	resized();
@@ -195,47 +195,81 @@ void RbmComponent::onDraw(DrawComponent &obj)
 {
 	if (&obj == DrawHidden)
 	{
-		downPass(toVisibleProbs(obj.getData()));
+		downPass(obj.getData());
 	}
 	if (&obj == DrawTraining)
 	{
-    	redrawReconstruction();
+		upDownPass(obj.getData());
 	}
 }
 
-void RbmComponent::redrawReconstruction()
+void RbmComponent::redrawReconstruction(const arma::mat& v)
 {
 	uint32_t i;
 
 	for (i=0; i < params().numGibbs; i++)
 	{
-		downPass(DrawReconstruction->getData());
+		DrawHidden->getData() = toHiddenProbs(v);
+		DrawHidden->DrawData();
+		DrawReconstruction->getData() = toVisibleProbs(DrawHidden->getData());
+		DrawReconstruction->DrawData();
 	}
-	DrawReconstruction->DrawData();
 }
 
 void RbmComponent::upPass(const arma::mat& v)
 {
-	DrawReconstruction->getData() = v;
+	DrawTraining->getData() = v;
+	DrawTraining->DrawData();
+	DrawHidden->getData() = toHiddenProbs(v);
+	DrawHidden->DrawData();
+	DrawReconstruction->getData() = toVisibleProbs(DrawHidden->getData());
 	DrawReconstruction->DrawData();
+	
 	if (next)
 	{
 		RbmComponent *pComp = static_cast<RbmComponent*> (next);
-		pComp->upPass(toHiddenProbs(v));
+		pComp->upPass(DrawHidden->getData());
 	}
 }
 
-void RbmComponent::downPass(const arma::mat& v)
+void RbmComponent::downPass(const arma::mat& h)
 {
-	DrawReconstruction->getData() = v;
+	DrawHidden->getData() = h;
+	DrawHidden->DrawData();
+	DrawReconstruction->getData() = toVisibleProbs(h);
 	DrawReconstruction->DrawData();
 	if (prev)
 	{
 		RbmComponent *pComp = static_cast<RbmComponent*> (prev);
-		pComp->downPass(prev->toVisibleProbs(v));
+		pComp->downPass(DrawReconstruction->getData());
 	}
 }
 
+void RbmComponent::upDownPass(const arma::mat& v)
+{
+	DrawTraining->getData() = v;
+	DrawTraining->DrawData();
+	DrawHidden->getData() = toHiddenProbs(v);
+	DrawHidden->DrawData();
+	DrawReconstruction->getData() = toVisibleProbs(DrawHidden->getData());
+	DrawReconstruction->DrawData();
+	
+	if (next)
+	{
+		RbmComponent *pComp = static_cast<RbmComponent*> (next);
+		pComp->upDownPass(DrawHidden->getData());
+	}
+	else if (prev)
+	{
+		if (prev)
+		{
+			RbmComponent *pComp = static_cast<RbmComponent*> (prev);
+			pComp->downPass(DrawReconstruction->getData());
+		}
+	}
+
+}
+
 arma::mat RbmComponent::getConvolutedWeight(arma::mat const &w)
 {
 //	if (next)
@@ -263,12 +297,7 @@ void RbmComponent::redrawWeights(size_t index)
 
 void RbmComponent::setTrainingData(arma::mat const& batch)
 {
-	
-	DrawTraining->getData() = trainingData(batch);
-	DrawTraining->DrawData();
-		
-	upPass(DrawTraining->getData());
-	redrawReconstruction();
+	upDownPass(trainingData(batch));
 }
 //[/MiscUserCode]
 
diff --git a/source/RbmComponent.hpp b/source/RbmComponent.hpp
index ae47245..ead4826 100644
--- a/source/RbmComponent.hpp
+++ b/source/RbmComponent.hpp
@@ -63,11 +63,11 @@ public:
 	void redrawWeights();
 	void redrawWeights(size_t index);
 
-	void redrawReconstruction();
+	void redrawReconstruction(arma::mat const &v);
 	
-	void downPass(arma::mat const &v);
 	void upPass(arma::mat const &v);
-	
+	void downPass(arma::mat const &v);
+	void upDownPass(arma::mat const &v);
 	arma::mat getConvolutedWeight(arma::mat const &h);
     ScopedPointer<DrawComponent> DrawTraining;
     ScopedPointer<DrawComponent> DrawHidden;